R/mutate.R

In duckdblabs/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(list(name = "mutate", x = .data, args = try_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_dots <- names(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]]
        name_dot <- names_dots[[i]]

        # Try expanding this `dot` if we see it is an `across()` call
        expanded <- duckplyr_expand_across(names_out, dot)

        if (is.null(expanded)) {
          # Nothing we can expand, create a list with just the 1 expression to
          # loop over
          quos <- set_names(list(dot), name_dot)
        } else {
          # Actually expanded an `across()` call, make sure to fix up names
          # again with the new set of dplyr quosures
          quos <- expanded
          quos <- fix_auto_name(quos)
        }

        names_quos <- names(quos)

        # Set up `exprs` outside the loop. All expressions expanded from an
        # `across()` are evaluated together using the same projection.
        exprs <- imap(set_names(names_out), relexpr_reference, rel = NULL)

        for (j in seq_along(quos)) {
          quo <- quos[[j]]
          new <- names_quos[[j]]

          names_new <- c(names_new, new)

          new_pos <- match(new, names_out, nomatch = length(names_out) + 1L)
          new_expr <- rel_translate(quo, names_data = names_out, alias = new, partition = by_names, need_window = TRUE)
          exprs[[new_pos]] <- new_expr

          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))
        }

        rel <- rel_project(rel, unname(exprs))
      }

      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
}