Skip to content
Snippets Groups Projects
Commit b9455d1f authored by adrian555's avatar adrian555 Committed by Shivaram Venkataraman
Browse files

[SPARK-11260][SPARKR] with() function support

Author: adrian555 <wzhuang@us.ibm.com>
Author: Adrian Zhuang <adrian555@users.noreply.github.com>

Closes #9443 from adrian555/with.
parent 8a5314ef
No related branches found
No related tags found
No related merge requests found
......@@ -83,6 +83,7 @@ exportMethods("arrange",
"unique",
"unpersist",
"where",
"with",
"withColumn",
"withColumnRenamed",
"write.df")
......
......@@ -2126,11 +2126,29 @@ setMethod("as.data.frame",
setMethod("attach",
signature(what = "DataFrame"),
function(what, pos = 2, name = deparse(substitute(what)), warn.conflicts = TRUE) {
cols <- columns(what)
stopifnot(length(cols) > 0)
newEnv <- new.env()
for (i in 1:length(cols)) {
assign(x = cols[i], value = what[, cols[i]], envir = newEnv)
}
newEnv <- assignNewEnv(what)
attach(newEnv, pos = pos, name = name, warn.conflicts = warn.conflicts)
})
#' Evaluate a R expression in an environment constructed from a DataFrame
#' with() allows access to columns of a DataFrame by simply referring to
#' their name. It appends every column of a DataFrame into a new
#' environment. Then, the given expression is evaluated in this new
#' environment.
#'
#' @rdname with
#' @title Evaluate a R expression in an environment constructed from a DataFrame
#' @param data (DataFrame) DataFrame to use for constructing an environment.
#' @param expr (expression) Expression to evaluate.
#' @param ... arguments to be passed to future methods.
#' @examples
#' \dontrun{
#' with(irisDf, nrow(Sepal_Width))
#' }
#' @seealso \link{attach}
setMethod("with",
signature(data = "DataFrame"),
function(data, expr, ...) {
newEnv <- assignNewEnv(data)
eval(substitute(expr), envir = newEnv, enclos = newEnv)
})
......@@ -1043,3 +1043,7 @@ setGeneric("as.data.frame")
#' @rdname attach
#' @export
setGeneric("attach")
#' @rdname with
#' @export
setGeneric("with")
......@@ -623,3 +623,16 @@ convertNamedListToEnv <- function(namedList) {
}
env
}
# Assign a new environment for attach() and with() methods
assignNewEnv <- function(data) {
stopifnot(class(data) == "DataFrame")
cols <- columns(data)
stopifnot(length(cols) > 0)
env <- new.env()
for (i in 1:length(cols)) {
assign(x = cols[i], value = data[, cols[i]], envir = env)
}
env
}
\ No newline at end of file
......@@ -1494,6 +1494,15 @@ test_that("attach() on a DataFrame", {
expect_error(age)
})
test_that("with() on a DataFrame", {
df <- createDataFrame(sqlContext, iris)
expect_error(Sepal_Length)
sum1 <- with(df, list(summary(Sepal_Length), summary(Sepal_Width)))
expect_equal(collect(sum1[[1]])[1, "Sepal_Length"], "150")
sum2 <- with(df, distinct(Sepal_Length))
expect_equal(nrow(sum2), 35)
})
unlink(parquetPath)
unlink(jsonPath)
unlink(jsonPathNa)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment