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: _wquantile
+    using StatsBase
+else
+    import ..CategoricalArrays: _wquantile
+    using ..StatsBase
+end
+
+_wquantile(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,17 @@ function find_breaks(v::AbstractVector, qs::AbstractVector)
     return breaks
 end
 
+# AbstractWeights method is defined in StatsBase extension
+# There is no in-place weighted quantile method in StatsBase
+_wquantile(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 +375,41 @@ 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 used when
+  computing quantiles (see `quantile` documentation in StatsBase).
 """
 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)
-    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"))
+    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
+        qs = quantile!(sorted_x, (1:(ngroups-1))/ngroups, sorted=true)
+    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)
+        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 = _wquantile(nm_x, nm_weights, (1:(ngroups-1))/ngroups)
     end
-    qs = quantile!(sorted_x, (1:(ngroups-1))/ngroups, sorted=true)
     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

@@ -1,6 +1,8 @@
 module TestExtras
 using Test
 using CategoricalArrays
+using StatsBase
+using Missings
 
 const ≅ = isequal
 
@@ -423,4 +425,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