R/uncount.R

In tidyfast: Fast Tidying of Data

#' Uncount
#'
#' Uncount a counted data table
#'
#' @param dt_ the data table to uncount
#' @param weights the counts for each
#' @param .remove should the weights variable be removed?
#' @param .id an optional new id variable, providing a unique id for each row
#'
#' @examples
#'
#' library(data.table)
#'
#' dt_count <- data.table(
#'   x = LETTERS[1:3],
#'   w = c(2, 1, 4)
#' )
#' uncount <- dt_uncount(dt_count, w, .id = "id")
#' uncount[] # note that `[]` forces the printing
#' @import data.table
#'
#' @return A data.table with a row for each uncounted column.
#'
#' @export
dt_uncount <- function(dt_, weights, .remove = TRUE, .id = NULL) {
  UseMethod("dt_uncount", dt_)
}

#' @export
dt_uncount.default <- function(dt_, weights, .remove = TRUE, .id = NULL) {
  if (isFALSE(is.data.table(dt_))) {
    dt_ <- as.data.table(dt_)
  }

  w <- substitute(weights)
  dt_ <- dt_[rep(1:.N, eval(w))]

  if (.remove) {
    dt_[, `:=`(paste(w), NULL)]
  }

  if (!is.null(.id)) {
    dt_[, `:=`(.id, 1:.N)]
  }

  dt_
}