Skip to content

Instantly share code, notes, and snippets.

@yutannihilation
Last active January 19, 2019 21:25
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save yutannihilation/958d2f2eb8b2fcddf3391a32a1740d6d to your computer and use it in GitHub Desktop.
Save yutannihilation/958d2f2eb8b2fcddf3391a32a1740d6d to your computer and use it in GitHub Desktop.
nest_column <- function(data, ..., .key = "data") {
  key_var <- rlang::as_string(rlang::ensym(.key))
  
  tie_vars <- unname(tidyselect::vars_select(names(data), ...))
  if (rlang::is_empty(tie_vars)) {
    tie_vars <- names(data)
  }
  
  if (dplyr::is_grouped_df(data)) {
    group_vars <- dplyr::group_vars(data)
  } else {
    group_vars <- setdiff(names(data), tie_vars)
  }
  tie_vars <- setdiff(tie_vars, group_vars)
  
  data <- dplyr::ungroup(data)
  if (rlang::is_empty(group_vars)) {
    return(tibble::tibble(!! key_var := data))
  }

  out <- dplyr::select(data, !!! rlang::syms(group_vars))
  out[[key_var]] <-   tied <- dplyr::select(data, !!! rlang::syms(tie_vars))
  out
}

unnest_column <- function(data, ...) {
  quos <- rlang::quos(...)
  if (rlang::is_empty(quos)) {
    list_cols <- names(data)[purrr::map_lgl(data, rlang::is_list)]
    quos <- rlang::syms(list_cols)
  }
  
  if (length(quos) == 0) {
    return(data)
  }
  
  tied <- as.list(dplyr::transmute(data, !!! quos))
  
  group_vars <- setdiff(names(data), names(tied))
  
  rest <- dplyr::select(data, !!!rlang::syms(group_vars))
  dplyr::bind_cols(rest, tied)
}

multi_spread <- function(data, key, ...) {
  data <- nest_column(data, ...)
  key <- rlang::enquo(key)
  
  df_var <- names(data)[purrr::map_lgl(data, is.data.frame)]
  rest_var <- setdiff(names(data), c(df_var, rlang::as_name(key)))

  data <- split(dplyr::select(data, -!!key), dplyr::pull(data, !!key))
  
  # TODO
  overall <- purrr::map(data, rest_var)
  overall <- sort(unique(unlist(overall)))

  data <- purrr::imap(data, function(d, nm) {
    res <- vctrs::vec_slice(d[[df_var]], match(overall, d[[rest_var]], nomatch = NA))
    names(res) <- paste(nm, names(res), sep = "_")
    res
  })
  
  dplyr::bind_cols(
    tibble::tibble(!! rest_var := overall),
    data
  )
}
data <- tibble::tribble(
  ~hw,   ~name,  ~mark,   ~pr,
  "hw1", "anna",    95,  "ok",
  "hw1", "alan",    90, "meh",
  "hw1", "carl",    85,  "ok",
  "hw2", "alan",    70, "meh",
  "hw2", "carl",    80,  "ok"
)

multi_spread(data, key = hw, mark, pr)
#> # A tibble: 3 x 5
#>   name  hw1_mark hw1_pr hw2_mark hw2_pr
#>   <chr>    <dbl> <chr>     <dbl> <chr> 
#> 1 alan        90 meh          70 meh   
#> 2 anna        95 ok           NA <NA>  
#> 3 carl        85 ok           80 ok
multi_spread(data, key = name, mark, pr)
#> # A tibble: 2 x 7
#>   hw    alan_mark alan_pr anna_mark anna_pr carl_mark carl_pr
#>   <chr>     <dbl> <chr>       <dbl> <chr>       <dbl> <chr>  
#> 1 hw1          90 meh            95 ok             85 ok     
#> 2 hw2          70 meh            NA <NA>           80 ok

Created on 2019-01-20 by the reprex package (v0.2.1)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment