Support weighted quantiles in cut

This requires adding an extension point for StatsBase.
Unfortunately more copies of the data and weights are done than necessary
as StatsBase does not support in-place weighted quantile! on pre-sorted data
nor taking a view of weights vectors (JuliaStats/StatsBase.jl#723).

Author: nalimilan

Project.toml

@@ -16,13 +16,15 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 
 [extensions]
 CategoricalArraysArrowExt = "Arrow"
 CategoricalArraysJSONExt = "JSON"
 CategoricalArraysRecipesBaseExt = "RecipesBase"
+CategoricalArraysStatsBaseExt = "StatsBase"
 CategoricalArraysSentinelArraysExt = "SentinelArrays"
 CategoricalArraysStructTypesExt = "StructTypes"
 
@@ -37,6 +39,7 @@ RecipesBase = "1.1"
 Requires = "1"
 SentinelArrays = "1"
 Statistics = "1"
+StatsBase = "0.15, 0.16, 0.17, 0.18, 0.19, 0.20, 0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
 StructTypes = "1"
 julia = "1.6"
 
@@ -49,8 +52,9 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Arrow", "Dates", "JSON", "JSON3", "Plots", "PooledArrays", "RecipesBase", "SentinelArrays", "StructTypes", "Test"]
+test = ["Arrow", "Dates", "JSON", "JSON3", "Plots", "PooledArrays", "RecipesBase", "SentinelArrays", "StatsBase", "StructTypes", "Test"]

ext/CategoricalArraysStatsBaseExt.jl

@@ -0,0 +1,13 @@
+module CategoricalArraysStatsBaseExt
+
+if isdefined(Base, :get_extension)
+    import CategoricalArrays: _quantile!
+    using StatsBase
+else
+    import ..CategoricalArrays: _quantile!
+    using ..StatsBase
+end
+
+_quantile!(x::AbstractArray, w::AbstractWeights, p::AbstractVector) = quantile(x, w, p)
+
+end

src/CategoricalArrays.jl

@@ -47,6 +47,7 @@ module CategoricalArrays
             @require JSON="682c06a0-de6a-54ab-a142-c8b1cf79cde6" include("../ext/CategoricalArraysJSONExt.jl")
             @require RecipesBase="3cdcf5f2-1ef4-517c-9805-6587b60abb01" include("../ext/CategoricalArraysRecipesBaseExt.jl")
             @require SentinelArrays="91c51154-3ec4-41a3-a24f-3f23e20d615c" include("../ext/CategoricalArraysSentinelArraysExt.jl")
+            @require StatsBase="2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" include("../ext/CategoricalArraysStatsBaseExt.jl")
             @require StructTypes="856f2bd8-1eba-4b0a-8007-ebc267875bd4" include("../ext/CategoricalArraysStructTypesExt.jl")
         end
     end

src/extras.jl

@@ -333,11 +333,19 @@ function find_breaks(v::AbstractVector, qs::AbstractVector)
     return breaks
 end
 
+_quantile!(x::AbstractArray, w::Nothing, p::AbstractVector) =
+    quantile!(x, p, sorted=true)
+# AbstractWeights method is defined in StatsBase extension
+# There is no in-place weighted quantile method in StatsBase
+_quantile!(x::AbstractArray, w::AbstractVector, p::AbstractVector) =
+    throw(ArgumentError("`weights` must be an `AbstractWeights` vector from StatsBase.jl"))
+
 """
     cut(x::AbstractArray, ngroups::Integer;
         labels::Union{AbstractVector{<:AbstractString},Function},
         sigdigits::Integer=3,
-        allowempty::Bool=false)
+        allowempty::Bool=false,
+        weights::Union{AbstractWeights, Nothing}=nothing)
 
 Cut a numeric array into `ngroups` quantiles.
 
@@ -369,19 +377,39 @@ quantiles.
   other than the last one are equal, generating empty intervals;
   when `true`, duplicate breaks are allowed and the intervals they generate are kept as
   unused levels (but duplicate labels are not allowed).
+* `weights::Union{AbstractWeights, Nothing}=nothing`: observations weights to pass to `quantile`.
 """
 function cut(x::AbstractArray, ngroups::Integer;
              labels::Union{AbstractVector{<:SupportedTypes},Function,Nothing}=nothing,
              sigdigits::Integer=3,
-             allowempty::Bool=false)
+             allowempty::Bool=false,
+             weights::Union{AbstractVector, Nothing}=nothing)
     ngroups >= 1 || throw(ArgumentError("ngroups must be strictly positive (got $ngroups)"))
-    sorted_x = eltype(x) >: Missing ? sort!(collect(skipmissing(x))) : sort(x)
+    if weights === nothing
+        sorted_x = eltype(x) >: Missing ? sort!(collect(skipmissing(x))) : sort(x)
+        min_x, max_x = first(sorted_x), last(sorted_x)
+        if (min_x isa Number && isnan(min_x)) ||
+            (max_x isa Number && isnan(max_x))
+            throw(ArgumentError("NaN values are not allowed in input vector"))
+        end
+    else
+        if eltype(x) >: Missing
+            nm_inds = findall(!ismissing, x)
+            nm_x = view(x, nm_inds)
+            # TODO: use a view once this is supported (JuliaStats/StatsBase.jl#723)
+            nm_weights = weights[nm_inds]
+        else
+            nm_x = x
+            nm_weights = weights
+        end
+        sorted_x = sort(nm_x)
+    end
     min_x, max_x = first(sorted_x), last(sorted_x)
     if (min_x isa Number && isnan(min_x)) ||
         (max_x isa Number && isnan(max_x))
         throw(ArgumentError("NaN values are not allowed in input vector"))
     end
-    qs = quantile!(sorted_x, (1:(ngroups-1))/ngroups, sorted=true)
+    qs = _quantile!(nm_x, nm_weights, (1:(ngroups-1))/ngroups)
     breaks = [min_x; find_breaks(sorted_x, qs); max_x]
     if !allowempty && !allunique(@view breaks[1:end-1])
         throw(ArgumentError("cannot compute $ngroups quantiles due to " *

test/15_extras.jl

@@ -423,4 +423,26 @@ end
 
 end
 
+@testset "cut with weighted quantiles" begin
+    @test_throws ArgumentError cut(1:3, 3, weights=1:3)
+
+    x = collect(Float64, 1:100)
+    w = fweights(repeat(1:10, inner=10))
+    y = cut(x, 10, weights=w)
+    @test levelcode.(y) == levelcode.(cut(x, quantile(x, w, (0:10)./10)))
+    @test levels(y) == ["[1, 29)", "[29, 43)", "[43, 53)", "[53, 62)", "[62, 70)",
+                        "[70, 77)", "[77, 83)", "[83, 89)", "[89, 95)", "[95, 100]"]
+
+    mx = allowmissing(x)
+    mx[2] = mx[10] = missing
+    nm_inds = .!ismissing.(mx)
+    y = cut(mx, 10, weights=w)
+    @test levelcode.(y) ≅ levelcode.(cut(mx, quantile(x[nm_inds], w[nm_inds], (0:10)./10)))
+    @test levels(y) == ["[1, 30)", "[30, 43)", "[43, 53)", "[53, 62)", "[62, 70)",
+                        "[70, 77)", "[77, 83)", "[83, 89)", "[89, 95)", "[95, 100]"]
+
+    x[5] = NaN
+    @test_throws ArgumentError cut(x, 3, weights=w)
+end
+
 end