R/transmute.R

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

# Generated by 02-duckplyr_df-methods.R
#' @rdname transmute.duckplyr_df
#' @export
transmute.duckplyr_df <- function(.data, ...) {
  force(.data)

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

  duckplyr_error <- rel_try(list(name = "transmute", x = .data, args = try_list(dots = enquos(...))),
    #' @section Fallbacks:
    #' There is no DuckDB translation in `transmute.duckplyr_df()`
    #' - with a selection that returns no columns:
    #'
    #' These features fall back to [dplyr::transmute()], see `vignette("fallback")` for details.
    "Zero-column result set not supported." = (length(dots) == 0),
    {
      rel <- duckdb_rel_from_df(.data)
      names_dots <- names(dots)
      names_new <- character()
      current_data <- rel_to_df(rel, prudence = "stingy")

      # Process expressions sequentially to allow referencing new variables
      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(current_data, dot)

        if (is.null(expanded)) {
          quos <- set_names(list(dot), name_dot)
        } else {
          quos <- expanded
          quos <- fix_auto_name(quos)
        }

        names_quos <- names(quos)

        # Set up exprs with all current columns
        exprs <- imap(set_names(names(current_data)), 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(current_data), nomatch = length(current_data) + j)
          new_expr <- rel_translate(
            quo,
            current_data,
            alias = new
          )
          exprs[[new_pos]] <- new_expr
        }

        rel <- rel_project(rel, unname(exprs))
        current_data <- rel_to_df(rel, prudence = "stingy")
      }

      # Select only the new columns (transmute behavior)
      final_exprs <- imap(set_names(names_new), relexpr_reference, rel = NULL)
      out_rel <- rel_project(rel, unname(final_exprs))
      out <- duckplyr_reconstruct(out_rel, .data)
      return(out)
    }
  )

  # dplyr forward
  check_prudence(.data, duckplyr_error)

  transmute <- dplyr$transmute.data.frame
  out <- transmute(.data, ...)
  return(out)

  # dplyr implementation
  dots <- check_transmute_args(...)
  dots <- dplyr_quosures(!!!dots)

  # We don't expose `.by` because `transmute()` is superseded
  by <- compute_by(by = NULL, data = .data)

  cols <- mutate_cols(.data, dots, by)

  out <- dplyr_col_modify(.data, cols)

  # Compact out `NULL` columns that got removed.
  # These won't exist in `out`, but we don't want them to look "new".
  # Note that `dplyr_col_modify()` makes it impossible to `NULL` a group column,
  # which we rely on below.
  cols <- compact_null(cols)

  # Retain expression columns in order of their appearance
  cols_expr <- names(cols)

  # Retain untouched group variables up front
  cols_group <- by$names
  cols_group <- setdiff(cols_group, cols_expr)

  cols_retain <- c(cols_group, cols_expr)

  dplyr_col_select(out, cols_retain)
}

duckplyr_transmute <- function(.data, ...) {
  try_fetch(
    .data <- as_duckplyr_df_impl(.data),
    error = function(e) {
      testthat::skip(conditionMessage(e))
    }
  )
  out <- transmute(.data, ...)
  class(out) <- setdiff(class(out), "duckplyr_df")
  out
}