R/count.R

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

# Generated by 02-duckplyr_df-methods.R
#' @rdname count.duckplyr_df
#' @export
count.duckplyr_df <- function(x, ..., wt = NULL, sort = FALSE, name = NULL, .drop = group_by_drop_default(x)) {
  force(x)

  dplyr_local_error_call()

  by <- dplyr_quosures(...)
  by <- fix_auto_name(by)

  by_exprs <- unname(map(by, quo_get_expr))
  is_name <- map_lgl(by_exprs, is_symbol)

  # For better error message
  if (!is.null(name)) {
    check_string(name, call = dplyr_error_call())
  }

  # Passing `name` reliably is surprisingly complicated.
  duckplyr_error <- rel_try(list(name = "count", x = x, args = try_list(dots = enquos(...), wt = enquo(wt), sort = sort, .drop = .drop)),
    #' @section Fallbacks:
    #' There is no DuckDB translation in `count.duckplyr_df()`
    #' - with complex expressions in `...`,
    #' - with `.drop = FALSE`,
    #' - with `sort = TRUE`.
    #'
    #' These features fall back to [dplyr::count()], see `vignette("fallback")` for details.
    "{.code count()} requires columns in {.arg ...}" = !all(is_name),
    "{.code count()} only implemented for {.arg .drop} = {.value TRUE}" = !.drop,
    "{.code count()} only implemented for {.arg sort} = {.value FALSE}" = sort,
    {
      rel <- duckdb_rel_from_df(x)

      by_chr <- map_chr(by_exprs, as_string)
      name_was_null <- is.null(name)
      name <- check_n_name(name, by_chr, call = dplyr_error_call())

      if (name %in% by_chr) {
        cli::cli_abort("Name clash in `count()`")
      }
      n <- tally_n(x, {{ wt }})

      groups <- rel_translate_dots(by, x)
      aggregates <- list(rel_translate(n, x, alias = name))

      out_rel <- rel_aggregate(rel, groups, unname(aggregates))
      if (length(groups) > 0) {
        sort_cols <- nexprs(names(groups))
        out_rel <- rel_order(out_rel, sort_cols)
      }

      out <- duckplyr_reconstruct(out_rel, x)
      return(out)
    }
  )

  # FIXME: optimize, no need to forward dots
  # out <- count(x_df, !!!quos, wt = {{ wt }}, sort = sort, name = name, .drop = .drop)

  # dplyr forward
  check_prudence(x, duckplyr_error)

  count <- dplyr$count.data.frame
  out <- count(x, ..., wt = {{ wt }}, sort = sort, name = name, .drop = .drop)
  return(out)

  # dplyr implementation
  dplyr_local_error_call()

  if (!missing(...)) {
    out <- group_by(x, ..., .add = TRUE, .drop = .drop)
  } else {
    out <- x
  }

  out <- tally(out, wt = !!enquo(wt), sort = sort, name = name)

  # Ensure grouping is transient
  out <- dplyr_reconstruct(out, x)

  out
}

duckplyr_count <- function(x, ...) {
  try_fetch(
    x <- as_duckplyr_df_impl(x),
    error = function(e) {
      testthat::skip(conditionMessage(e))
    }
  )
  out <- count(x, ...)
  class(out) <- setdiff(class(out), "duckplyr_df")
  out
}