Support different types in colQuantiles

All types expect for type = 7L (ie. the default) are calculated
very inefficiently, but expanding each column to a dense vector

This fixes the last open bit of #7

Author: const-ae

R/methods.R

@@ -507,14 +507,27 @@ setMethod("colCollapse", signature(x = "xgCMatrix"),
 #' @inherit MatrixGenerics::colQuantiles
 #' @export
 setMethod("colQuantiles", signature(x = "xgCMatrix"),
-          function(x, rows = NULL, cols = NULL, probs = seq(from = 0, to = 1, by = 0.25), na.rm=FALSE, drop = TRUE){
+          function(x, rows = NULL, cols = NULL, probs = seq(from = 0, to = 1, by = 0.25), na.rm=FALSE, type = 7L, drop = TRUE){
   if(! is.null(rows)){
     x <- x[rows, , drop = FALSE]
   }
   if(! is.null(cols)){
     x <- x[, cols, drop = FALSE]
   }
-  mat <- dgCMatrix_colQuantiles(x, probs, na_rm = na.rm)
+  if(type == 7L){
+    mat <- dgCMatrix_colQuantiles(x, probs, na_rm = na.rm)
+  }else{
+    mat <- t(expand_and_reduce_sparse_matrix_to_matrix(x, n_result_rows = length(probs), function(values){
+      if(na.rm){
+        values <- values[!is.na(values)]
+        stats::quantile(values, probs = probs, na.rm = na.rm, names = FALSE, type = type)
+      }else if(any(is.na(values))){
+        rep(NA_real_, length(probs))
+      }else{
+        stats::quantile(values, probs = probs, na.rm = na.rm, names = FALSE, type = type)
+      }
+    }))
+  }
   # Add dim names
   digits <- max(2L, getOption("digits"))
   colnames(mat) <- sprintf("%.*g%%", digits, 100 * probs)

R/sparse_matrix_iterator.R

@@ -34,3 +34,26 @@ reduce_sparse_matrix_to_matrix <- function(sp_mat, n_result_rows, reduce_functio
     }
   }
 }
+
+
+expand_and_reduce_sparse_matrix_to_matrix <- function(sp_mat, n_result_rows, reduce_function = function(dense_values){ NA_real_}){
+  if(length(sp_mat@p) == 0){
+    numeric(0)
+  }else{
+    res <- vapply(seq_len(length(sp_mat@p)-1), function(index){
+      start_pos <- sp_mat@p[index]
+      end_pos <- sp_mat@p[index + 1]
+      number_of_zeros <- nrow(sp_mat) - (end_pos - start_pos)
+      values <- sp_mat@x[start_pos + seq_len(end_pos - start_pos)]
+      row_indices <- sp_mat@i[start_pos + seq_len(end_pos - start_pos)]
+      dense_values <- rep(0, nrow(sp_mat))
+      dense_values[row_indices + 1] <- values
+      reduce_function(dense_values)
+    }, FUN.VALUE = rep(0.0, n_result_rows))
+    if(n_result_rows == 1){
+      matrix(res, nrow=1, ncol=length(res))
+    }else{
+      res
+    }
+  }
+}

man/colQuantiles-xgCMatrix-method.Rd

@@ -199,6 +199,17 @@ for(idx in seq_along(matrix_list)){
     expect_equal(colQuantiles(sp_mat), matrixStats::colQuantiles(mat))
     expect_equal(colQuantiles(sp_mat, na.rm=TRUE), matrixStats::colQuantiles(mat, na.rm=TRUE))
     expect_equal(colQuantiles(sp_mat, rows = row_subset, cols = col_subset), matrixStats::colQuantiles(mat, rows = row_subset, cols = col_subset))
+
+    expect_equal(colQuantiles(sp_mat, type = 1L), matrixStats::colQuantiles(mat, type = 1L))
+    expect_equal(colQuantiles(sp_mat, type = 2L), matrixStats::colQuantiles(mat, type = 2L))
+    expect_equal(colQuantiles(sp_mat, type = 3L), matrixStats::colQuantiles(mat, type = 3L))
+
+    expect_equal(colQuantiles(sp_mat, type = 4L), matrixStats::colQuantiles(mat, type = 4L))
+    expect_equal(colQuantiles(sp_mat, type = 5L), matrixStats::colQuantiles(mat, type = 5L))
+    expect_equal(colQuantiles(sp_mat, type = 6L), matrixStats::colQuantiles(mat, type = 6L))
+    expect_equal(colQuantiles(sp_mat, type = 7L), matrixStats::colQuantiles(mat, type = 7L))
+    expect_equal(colQuantiles(sp_mat, type = 8L), matrixStats::colQuantiles(mat, type = 8L))
+    expect_equal(colQuantiles(sp_mat, type = 9L), matrixStats::colQuantiles(mat, type = 9L))
   })