Add wts argument to sparse_mean()

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: sparsevctrs
 Title: Sparse Vectors for Use in Data Frames
-Version: 0.2.0.9003
+Version: 0.2.0.9004
 Authors@R: c(
     person("Emil", "Hvitfeldt", , "emil.hvitfeldt@posit.co", role = c("aut", "cre"),
            comment = c(ORCID = "0000-0002-0679-1945")),

NEWS.md

@@ -10,6 +10,8 @@
 
 * Adding the arithmatic function `sparse_multiplication()`. (#93)
 
+* Adding `wts` argument to `sparse_mean()`. (#95)
+
 # sparsevctrs 0.2.0
 
 ## New Functions

R/sparse_mean.R

@@ -1,8 +1,10 @@
 #' Calculate mean from sparse vectors
 #'
 #' @param x A sparse numeric vector.
+#' @param wts A numeric vector, should be same length as `x`.
 #' @param na_rm Logical, whether to remove missing values. Defaults to `FALSE`.
 #'
+#'
 #' @details
 #' This function, as with any of the other helper functions assumes that the
 #' input `x` is a sparse numeric vector. This is done for performance reasons,
@@ -33,7 +35,11 @@
 #' )
 #'
 #' @export
-sparse_mean <- function(x, na_rm = FALSE) {
+sparse_mean <- function(x, wts = NULL, na_rm = FALSE) {
+  if (!is.null(wts)) {
+    x <- sparse_multiplication(x, wts)
+  }
+
   default <- sparse_default(x)
   values <- sparse_values(x)
   len_values <- length(values)
@@ -46,15 +52,23 @@ sparse_mean <- function(x, na_rm = FALSE) {
 
   res <- sum(values, na.rm = na_rm)
 
-  if (default != 0) {
+  if (!is.na(default) && default != 0) {
     res <- res + (x_len - len_values) * default
   }
 
   if (na_rm) {
     x_len <- x_len - sum(is.na(values))
   }
 
-  res <- res / x_len
+  if (is.null(wts)) {
+    res <- res / x_len
+  } else {
+    na_loc <- sparse_which_na(x)
+    if (length(na_loc) > 0) {
+      wts <- wts[-na_loc]
+    }
+    res <- res / sum(wts)
+  }
 
   res
 }

tests/testthat/test-sparse_mean.R

@@ -35,3 +35,44 @@ test_that("sparse_mean() works", {
 
   expect_equal(mean(x), sparse_mean(x))
 })
+
+test_that("sparse_mean() works with wts argument", {
+  x <- sparse_double(10, 5, 1000)
+  wts <- (1:1000)[]
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+
+  x <- sparse_double(c(10, -10), c(5, 100), 1000)
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+
+  x <- sparse_double(c(NA, 10, 30), 1:3, 1000)
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+
+  x <- sparse_double(c(NA, 10, 30), 1:3, 1000, default = 100)
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+
+  x <- sparse_double(c(NA, 10, 30), 1:3, 1000)
+
+  expect_equal(
+    weighted.mean(x, wts, na.rm = TRUE),
+    sparse_mean(x, wts = wts, na_rm = TRUE)
+  )
+
+  x <- sparse_double(c(NA, 10, 30), 1:3, 1000, default = 100)
+
+  expect_equal(
+    weighted.mean(x, wts, na.rm = TRUE),
+    sparse_mean(x, wts = wts, na_rm = TRUE)
+  )
+
+  x <- sparse_double(numeric(), integer(), 1000)
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+
+  x <- sparse_double(numeric(), integer(), 1000, default = 100)
+
+  expect_equal(weighted.mean(x, wts), sparse_mean(x, wts = wts))
+})