R/foxtrot.R

In romainfrancois/dance: Tibble Dancing

#' @rdname jive
#' @export
foxtrot <- function(.tbl, ..., .env = caller_env()) {
  # evaluate all the formulas in each group
  c(., steps) %<-% ballet(.tbl, ..., .env = .env)

  map(steps, ~as_tibble_splice(.))
}

#' @rdname jive
#' @export
bachata <- function(.tbl, ..., .name = "data", .env = caller_env()) {
  vec_cbind(polka(.tbl), !!.name := foxtrot(.tbl, ..., .env = .env))
}

#' Modify
#'
#' Applies the [ballet()] with no restriction on the [vec_size()] of the outputs.
#'
#' @param .tbl A data frame, most likely a grouped data frame
#' @param ...,.env formulas for each column to create, and parent environment, see [ballet()]
#' @param .name Name of the packed column made by `bachata()`
#'
#' @return
#'
#' - `foxtrot()` returns a list of tibbles, one per group of `.tbl`. This is
#'   close in functionality to [dplyr::group_map()].
#'
#' - `bachata()` returns a tibble that [vctrs::vec_cbind()] the results of [polka()] and
#'   the result from `foxtrot()` as a list column
#'
#' - `jive()` row binds the results from `foxtrot()` and reconstructs the grouping structure.
#'   This is similar in functionality to `dplyr::group_modify()`
#'
#' @examples
#' g <- group_by(iris, Species)
#'
#' \dontrun{
#' if (requireNamespace("broom", quietly = TRUE)) {
#'   ### foxtrot() to return a list of tibbles
#'
#'   # A list of tibbles with the `quantiles` column derived from the
#'   # formula on each group
#'   g %>%
#'     foxtrot(
#'       quantiles = ~quantile(Petal.Length, probs = c(0.25, 0.5, 0.75))
#'     )
#'
#'   g %>%
#'     foxtrot(
#'       ~ broom::tidy(lm(Petal.Length ~ Sepal.Length))
#'     )
#'
#'   ### bachata() to nest the results of foxtrot()
#'
#'   # bachata() = polka() + a list column of 2 columns tibbles
#'   g %>%
#'     bachata(
#'        Sepal.Length = ~quantile(Sepal.Length, probs = c(0.25, 0.5, 0.75)),
#'        Sepal.Width  = ~quantile(Sepal.Width, probs = c(0.25, 0.5, 0.75)),
#'     )
#'
#'   # because the formula is unnamed, the columns of the tibbles
#'   # generated by broom::tidy() are auto-spliced
#'   g %>%
#'     bachata(
#'       ~ broom::tidy(lm(Petal.Length ~ Sepal.Length))
#'     )
#'
#'   ### jive()
#'   g %>%
#'     jive(
#'        Sepal.Length = ~quantile(Sepal.Length, probs = c(0.25, 0.5, 0.75)),
#'        Sepal.Width  = ~quantile(Sepal.Width, probs = c(0.25, 0.5, 0.75)),
#'     )
#'
#'   # unnamed -> columns are spliced
#'   g %>%
#'     jive(
#'       ~ broom::tidy(lm(Petal.Length ~ Sepal.Length))
#'     )
#'
#'   # named -> packed column
#'   g %>%
#'     jive(
#'       model = ~ broom::tidy(lm(Petal.Length ~ Sepal.Length))
#'     )
#' }
#' }
#' @export
jive <- function(.tbl, ..., .env = caller_env()) {
  chunks <- foxtrot(.tbl, ..., .env = .env)
  sizes  <- map_int(chunks, nrow)
  keys   <- group_keys(.tbl)
  gps    <- groups(.tbl)

  out <- vec_cbind(
    keys[rep(seq_len(nrow(keys)), sizes), ],
    vec_rbind(!!!chunks)
  )
  group_by(out, !!!gps)
}