Skip to content

Commit

Permalink
Pass caller environment as argument to translate() functions (#45)
Browse files Browse the repository at this point in the history
* pass caller environment as argument to translate() functions

* same thing for summarize()
  • Loading branch information
etiennebacher authored Sep 17, 2023
1 parent 5059946 commit 1459111
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 71 deletions.
2 changes: 1 addition & 1 deletion R/mutate.R
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ pl_mutate <- function(.data, ...) {
is_grouped <- !is.null(grps)
to_drop <- list()

polars_exprs <- translate_dots(.data = .data, ...)
polars_exprs <- translate_dots(.data = .data, ..., env = rlang::caller_env())

for (i in seq_along(polars_exprs)) {
sub <- polars_exprs[[i]]
Expand Down
106 changes: 53 additions & 53 deletions R/summarize.R
Original file line number Diff line number Diff line change
@@ -1,53 +1,53 @@
#' Summarize each group down to one row
#'
#' `pl_summarize()` returns one row for each combination of grouping variables
#' (one difference with `dplyr::summarize()` is that `pl_summarize()` only
#' accepts grouped data). It will contain one column for each grouping variable
#' and one column for each of the summary statistics that you have specified.
#'
#' @param .data A Polars Data/LazyFrame
#' @inheritParams pl_mutate
#'
#' @export
#' @examples
#' mtcars |>
#' as_polars() |>
#' pl_group_by(cyl) |>
#' pl_summarize(gear = mean(gear), gear2 = sd(gear))


pl_summarize <- function(.data, ...) {

check_polars_data(.data)

grps <- attributes(.data)$pl_grps
mo <- attributes(.data)$maintain_grp_order
if (is.null(mo)) mo <- FALSE
is_grouped <- !is.null(grps)

if (!is_grouped) {
rlang::abort("`pl_summarize()` only works on grouped data.")
}

polars_exprs <- translate_dots(.data = .data, ...)

for (i in seq_along(polars_exprs)) {
sub <- polars_exprs[[i]]
to_drop <- names(empty_elems(sub))
sub <- compact(sub)

if (length(sub) > 0) {
.data <- .data$groupby(grps, maintain_order = mo)$agg(sub)
}

if (length(to_drop) > 0) {
.data <- .data$drop(to_drop)
}
}

.data
}

#' @rdname pl_summarize
#' @export
pl_summarise <- pl_summarize
#' Summarize each group down to one row
#'
#' `pl_summarize()` returns one row for each combination of grouping variables
#' (one difference with `dplyr::summarize()` is that `pl_summarize()` only
#' accepts grouped data). It will contain one column for each grouping variable
#' and one column for each of the summary statistics that you have specified.
#'
#' @param .data A Polars Data/LazyFrame
#' @inheritParams pl_mutate
#'
#' @export
#' @examples
#' mtcars |>
#' as_polars() |>
#' pl_group_by(cyl) |>
#' pl_summarize(gear = mean(gear), gear2 = sd(gear))


pl_summarize <- function(.data, ...) {

check_polars_data(.data)

grps <- attributes(.data)$pl_grps
mo <- attributes(.data)$maintain_grp_order
if (is.null(mo)) mo <- FALSE
is_grouped <- !is.null(grps)

if (!is_grouped) {
rlang::abort("`pl_summarize()` only works on grouped data.")
}

polars_exprs <- translate_dots(.data = .data, ..., env = rlang::caller_env())

for (i in seq_along(polars_exprs)) {
sub <- polars_exprs[[i]]
to_drop <- names(empty_elems(sub))
sub <- compact(sub)

if (length(sub) > 0) {
.data <- .data$groupby(grps, maintain_order = mo)$agg(sub)
}

if (length(to_drop) > 0) {
.data <- .data$drop(to_drop)
}
}

.data
}

#' @rdname pl_summarize
#' @export
pl_summarise <- pl_summarize
34 changes: 17 additions & 17 deletions R/utils-expr.R
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
#' @import rlang

translate_dots <- function(.data, ...) {
translate_dots <- function(.data, ..., env) {
dots <- enexprs(...)
new_vars <- c()
out <- lapply(seq_along(dots), \(x) {
tmp <- translate_expr(.data = .data, dots[[x]], new_vars)
tmp <- translate_expr(.data = .data, dots[[x]], new_vars, env = env)
new_vars <<- c(new_vars, names(dots)[x])
tmp
})
Expand All @@ -22,7 +22,7 @@ translate_dots <- function(.data, ...) {
out
}

translate_expr <- function(.data, quo, new_vars) {
translate_expr <- function(.data, quo, new_vars, env) {

names_data <- pl_colnames(.data)

Expand Down Expand Up @@ -51,7 +51,7 @@ translate_expr <- function(.data, quo, new_vars) {
expr <- unpack_across(.data, expr)
}

translate <- function(expr) {
translate <- function(expr, env) {

# prepare function and arg if the user provided an anonymous function in
# across()
Expand Down Expand Up @@ -85,7 +85,7 @@ translate_expr <- function(.data, quo, new_vars) {
if (expr_char %in% names_data || expr_char %in% unlist(new_vars)) {
polars_col(expr_char)
} else {
val <- eval_tidy(expr, env = caller_env(3))
val <- eval_tidy(expr, env = env)
polars_constant(val)
}
},
Expand Down Expand Up @@ -118,8 +118,8 @@ translate_expr <- function(.data, quo, new_vars) {
"%in%" = {
out <- tryCatch(
{
lhs <- translate(expr[[2]])
rhs <- translate(expr[[3]])
lhs <- translate(expr[[2]], env = env)
rhs <- translate(expr[[3]], env = env)
if (is.list(rhs)) {
rhs <- unlist(rhs)
}
Expand All @@ -138,7 +138,7 @@ translate_expr <- function(.data, quo, new_vars) {
"is.na" = {
out <- tryCatch(
{
inside <- translate(expr[[2]])
inside <- translate(expr[[2]], env = env)
inside$is_null()
},
error = identity
Expand All @@ -148,7 +148,7 @@ translate_expr <- function(.data, quo, new_vars) {
"is.nan" = {
out <- tryCatch(
{
inside <- translate(expr[[2]])
inside <- translate(expr[[2]], env = env)
inside$is_nan()
},
error = identity
Expand All @@ -169,7 +169,7 @@ translate_expr <- function(.data, quo, new_vars) {
if (startsWith(obj_name, ".__tidypolars__across_fn")) {
fn <- eval_bare(global_env()[[obj_name]])
col_name <- sym(col_name)
args <- translate(col_name)
args <- translate(col_name, env = env)
suppressWarnings({
tr <- try(do.call(fn, list(args)), silent = TRUE)
})
Expand All @@ -179,15 +179,15 @@ translate_expr <- function(.data, quo, new_vars) {
abort(
c("Could not evaluate an anonymous function in `across()`.",
"i" = "Are you sure the anonymous function returns a Polars expression?"),
call = caller_env(7)
call = env
)
}
} else {
abort(paste("Unknown function:", name), call = caller_env(5))
abort(paste("Unknown function:", name), call = env)
}
}

args <- lapply(as.list(expr[-1]), translate)
args <- lapply(as.list(expr[-1]), translate, env = env)
if (name %in% known_functions) {
name <- r_polars_funs$polars_funs[r_polars_funs$r_funs == name][1]
name <- paste0("pl_", name)
Expand All @@ -203,24 +203,24 @@ translate_expr <- function(.data, quo, new_vars) {

)
} else {
abort(e$message, call = caller_env(9))
abort(e$message, call = env)
}
}
)
},

abort(
paste("Internal: Unknown type", typeof(expr)),
call = caller_env(4)
call = env
)
)
}

# happens because across() calls get split earlier
if ((is.vector(expr) && length(expr) > 1) || is.list(expr)) {
lapply(expr, translate)
lapply(expr, translate, env = env)
} else {
translate(expr)
translate(expr, env = env)
}
}

Expand Down

0 comments on commit 1459111

Please sign in to comment.