/R/pivot_longer.R

https://github.com/TysonStanley/tidyfast · R · 97 lines · 47 code · 12 blank · 38 comment · 8 complexity · d345d1eb2ee0f6e9a924f5658aad3013 MD5 · raw file

  1. #' Pivot data from wide to long
  2. #'
  3. #'
  4. #' \code{dt_pivot_wider()} "widens" data, increasing the number of columns and
  5. #' decreasing the number of rows. The inverse transformation is
  6. #' \code{dt_pivot_longer()}. Syntax based on the \code{tidyr} equivalents.
  7. #'
  8. #' @param dt_ The data table to pivot longer
  9. #' @param cols Column selection. If empty, uses all columns. Can use -colname to unselect column(s)
  10. #' @param names_to Name of the new "names" column. Must be a string.
  11. #' @param values_to Name of the new "values" column. Must be a string.
  12. #' @param values_drop_na If TRUE, rows will be dropped that contain NAs.
  13. #' @param ... Additional arguments to pass to `melt.data.table()`
  14. #'
  15. #' @return A reshaped data.table into longer format
  16. #'
  17. #' @examples
  18. #'
  19. #' library(data.table)
  20. #' example_dt <- data.table(x = c(1, 2, 3), y = c(4, 5, 6), z = c("a", "b", "c"))
  21. #'
  22. #' dt_pivot_longer(example_dt,
  23. #' cols = c(x, y),
  24. #' names_to = "stuff",
  25. #' values_to = "things"
  26. #' )
  27. #'
  28. #' dt_pivot_longer(example_dt,
  29. #' cols = -z,
  30. #' names_to = "stuff",
  31. #' values_to = "things"
  32. #' )
  33. #' @importFrom data.table melt
  34. #' @importFrom stats setNames
  35. #'
  36. #' @export
  37. dt_pivot_longer <- function(dt_,
  38. cols = NULL,
  39. names_to = "name",
  40. values_to = "value",
  41. values_drop_na = FALSE,
  42. ...) {
  43. UseMethod("dt_pivot_longer", dt_)
  44. }
  45. #' @export
  46. dt_pivot_longer.default <- function(dt_,
  47. cols = NULL,
  48. names_to = "name",
  49. values_to = "value",
  50. values_drop_na = FALSE,
  51. ...) {
  52. if (!is.data.frame(dt_)) stop("dt_ must be a data.frame or data.table")
  53. if (!is.data.table(dt_)) dt_ <- as.data.table(dt_)
  54. names <- colnames(dt_)
  55. if (is.null(substitute(cols))) {
  56. # All columns if cols = NULL
  57. cols <- names
  58. } else {
  59. cols <- column_selector(dt_, substitute(c(cols)))
  60. }
  61. if (length(cols) == 0) warning("No columns remaining after removing")
  62. id_vars <- names[!names %in% cols]
  63. melt(
  64. data = dt_,
  65. id.vars = id_vars,
  66. measure.vars = cols,
  67. variable.name = names_to,
  68. value.name = values_to,
  69. ...,
  70. na.rm = values_drop_na,
  71. variable.factor = FALSE,
  72. value.factor = FALSE
  73. )
  74. }
  75. column_selector <- function(.data, select_vars) {
  76. data_names <- colnames(.data)
  77. data_vars <- setNames(as.list(seq_along(.data)), data_names)
  78. select_index <- eval(select_vars, data_vars)
  79. keep_index <- unique(select_index[select_index > 0])
  80. if (length(keep_index) == 0) keep_index <- seq_along(.data)
  81. drop_index <- unique(abs(select_index[select_index < 0]))
  82. select_index <- setdiff(keep_index, drop_index)
  83. select_vars <- data_names[select_index]
  84. select_vars
  85. }