R/mutate.R

In duckplyr: A 'DuckDB'-Backed Version of 'dplyr'

# Generated by 02-duckplyr_df-methods.R
#' @export
mutate.duckplyr_df <- function(.data, ..., .by = NULL, .keep = c("all", "used", "unused", "none"), .before = NULL, .after = NULL) {
  by_arg <- enquo(.by)
  keep <- arg_match(.keep)

  by_names <- eval_select_by(by_arg, .data)

  # Our implementation
  rel_try(call = list(name = "mutate", x = .data, args = list(dots = enquos(...), .by = by_arg, .keep = .keep)),
    "Implemented for all cases?" = FALSE,
    {
      rel <- duckdb_rel_from_df(.data)

      if (length(by_names) > 0) {
        rel <- oo_prep(rel)
      }

      dots <- dplyr_quosures(...)
      dots <- fix_auto_name(dots)

      names_used <- character()
      names_new <- character()
      names_out <- rel_names(rel)

      # FIXME: use fewer projections
      for (i in seq_along(dots)) {
        dot <- dots[[i]]

        new <- names(dots)[[i]]

        names_new <- c(names_new, new)

        new_pos <- match(new, names_out, nomatch = length(names_out) + 1L)
        exprs <- imap(set_names(names_out), relexpr_reference, rel = NULL)
        new_expr <- rel_translate(dot, names_data = names_out, alias = new, partition = by_names, need_window = TRUE)
        exprs[[new_pos]] <- new_expr

        rel <- rel_project(rel, unname(exprs))
        names_out[[new_pos]] <- new

        new_names_used <- intersect(attr(new_expr, "used"), names(.data))
        names_used <- c(names_used, setdiff(new_names_used, names_used))
      }

      if (length(by_names) > 0) {
        rel <- oo_restore(rel)
      }

      out <- rel_to_df(rel)

      out <- dplyr_reconstruct(out, .data)

      names_original <- names(.data)

      out <- mutate_relocate(
        out = out,
        before = {{ .before }},
        after = {{ .after }},
        names_original = names_original
      )

      used <- set_names(names(out) %in% names_used, names(out))
      names_groups <- by_names

      out <- duckplyr_mutate_keep(
        out = out,
        keep = keep,
        used = used,
        names_new = names_new,
        names_groups = names_groups
      )

      return(out)
    }
  )

  # dplyr forward
  mutate <- dplyr$mutate.data.frame
  out <- mutate(.data, ..., .by = {{ .by }}, .keep = .keep, .before = {{ .before }}, .after = {{ .after }})
  return(out)

  # dplyr implementation
  keep <- arg_match0(.keep, values = c("all", "used", "unused", "none"))

  by <- compute_by({{ .by }}, .data, by_arg = ".by", data_arg = ".data")

  cols <- mutate_cols(.data, dplyr_quosures(...), by)
  used <- attr(cols, "used")

  out <- dplyr_col_modify(.data, cols)

  names_original <- names(.data)

  out <- mutate_relocate(
    out = out,
    before = {{ .before }},
    after = {{ .after }},
    names_original = names_original
  )

  names_new <- names(cols)
  names_groups <- by$names

  out <- mutate_keep(
    out = out,
    keep = keep,
    used = used,
    names_new = names_new,
    names_groups = names_groups
  )

  out
}

duckplyr_mutate <- function(.data, ...) {
  try_fetch(
    .data <- as_duckplyr_df(.data),
    error = function(e) {
      testthat::skip(conditionMessage(e))
    }
  )
  out <- mutate(.data, ...)
  class(out) <- setdiff(class(out), "duckplyr_df")
  out
}