Repository: spark Updated Branches: refs/heads/branch-2.1 1022049c7 -> 82fcc1330
[SPARK-19130][SPARKR] Support setting literal value as column implicitly ## What changes were proposed in this pull request? ``` df$foo <- 1 ``` instead of ``` df$foo <- lit(1) ``` ## How was this patch tested? unit tests Author: Felix Cheung <felixcheun...@hotmail.com> Closes #16510 from felixcheung/rlitcol. (cherry picked from commit d749c06677c2fd383733337f1c00f542da122b8d) Signed-off-by: Shivaram Venkataraman <shiva...@cs.berkeley.edu> Project: http://git-wip-us.apache.org/repos/asf/spark/repo Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/82fcc133 Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/82fcc133 Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/82fcc133 Branch: refs/heads/branch-2.1 Commit: 82fcc133040cb5ef32f10df73fcb9fd8914aa9c1 Parents: 1022049 Author: Felix Cheung <felixcheun...@hotmail.com> Authored: Wed Jan 11 08:29:09 2017 -0800 Committer: Shivaram Venkataraman <shiva...@cs.berkeley.edu> Committed: Wed Jan 11 08:29:30 2017 -0800 ---------------------------------------------------------------------- R/pkg/R/DataFrame.R | 22 +++++++++++++++++----- R/pkg/R/utils.R | 4 ++++ R/pkg/inst/tests/testthat/test_sparkSQL.R | 18 ++++++++++++++++++ 3 files changed, 39 insertions(+), 5 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/spark/blob/82fcc133/R/pkg/R/DataFrame.R ---------------------------------------------------------------------- diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R index 058a77e..c79b1d3 100644 --- a/R/pkg/R/DataFrame.R +++ b/R/pkg/R/DataFrame.R @@ -1721,14 +1721,21 @@ setMethod("$", signature(x = "SparkDataFrame"), getColumn(x, name) }) -#' @param value a Column or \code{NULL}. If \code{NULL}, the specified Column is dropped. +#' @param value a Column or an atomic vector in the length of 1 as literal value, or \code{NULL}. +#' If \code{NULL}, the specified Column is dropped. #' @rdname select #' @name $<- #' @aliases $<-,SparkDataFrame-method #' @note $<- since 1.4.0 setMethod("$<-", signature(x = "SparkDataFrame"), function(x, name, value) { - stopifnot(class(value) == "Column" || is.null(value)) + if (class(value) != "Column" && !is.null(value)) { + if (isAtomicLengthOne(value)) { + value <- lit(value) + } else { + stop("value must be a Column, literal value as atomic in length of 1, or NULL") + } + } if (is.null(value)) { nx <- drop(x, name) @@ -1941,10 +1948,10 @@ setMethod("selectExpr", #' #' @param x a SparkDataFrame. #' @param colName a column name. -#' @param col a Column expression. +#' @param col a Column expression, or an atomic vector in the length of 1 as literal value. #' @return A SparkDataFrame with the new column added or the existing column replaced. #' @family SparkDataFrame functions -#' @aliases withColumn,SparkDataFrame,character,Column-method +#' @aliases withColumn,SparkDataFrame,character-method #' @rdname withColumn #' @name withColumn #' @seealso \link{rename} \link{mutate} @@ -1957,11 +1964,16 @@ setMethod("selectExpr", #' newDF <- withColumn(df, "newCol", df$col1 * 5) #' # Replace an existing column #' newDF2 <- withColumn(newDF, "newCol", newDF$col1) +#' newDF3 <- withColumn(newDF, "newCol", 42) #' } #' @note withColumn since 1.4.0 setMethod("withColumn", - signature(x = "SparkDataFrame", colName = "character", col = "Column"), + signature(x = "SparkDataFrame", colName = "character"), function(x, colName, col) { + if (class(col) != "Column") { + if (!isAtomicLengthOne(col)) stop("Literal value must be atomic in length of 1") + col <- lit(col) + } sdf <- callJMethod(x@sdf, "withColumn", colName, col@jc) dataFrame(sdf) }) http://git-wip-us.apache.org/repos/asf/spark/blob/82fcc133/R/pkg/R/utils.R ---------------------------------------------------------------------- diff --git a/R/pkg/R/utils.R b/R/pkg/R/utils.R index 1283449..74b3e50 100644 --- a/R/pkg/R/utils.R +++ b/R/pkg/R/utils.R @@ -863,3 +863,7 @@ basenameSansExtFromUrl <- function(url) { # then, strip extension by the last '.' sub("([^.]+)\\.[[:alnum:]]+$", "\\1", filename) } + +isAtomicLengthOne <- function(x) { + is.atomic(x) && length(x) == 1 +} http://git-wip-us.apache.org/repos/asf/spark/blob/82fcc133/R/pkg/inst/tests/testthat/test_sparkSQL.R ---------------------------------------------------------------------- diff --git a/R/pkg/inst/tests/testthat/test_sparkSQL.R b/R/pkg/inst/tests/testthat/test_sparkSQL.R index 4490f31..0be924f 100644 --- a/R/pkg/inst/tests/testthat/test_sparkSQL.R +++ b/R/pkg/inst/tests/testthat/test_sparkSQL.R @@ -1001,6 +1001,17 @@ test_that("select operators", { expect_equal(columns(df), c("name", "age", "age2")) expect_equal(count(where(df, df$age2 == df$age * 2)), 2) + df$age2 <- 21 + expect_equal(columns(df), c("name", "age", "age2")) + expect_equal(count(where(df, df$age2 == 21)), 3) + + df$age2 <- c(22) + expect_equal(columns(df), c("name", "age", "age2")) + expect_equal(count(where(df, df$age2 == 22)), 3) + + expect_error(df$age3 <- c(22, NA), + "value must be a Column, literal value as atomic in length of 1, or NULL") + # Test parameter drop expect_equal(class(df[, 1]) == "SparkDataFrame", T) expect_equal(class(df[, 1, drop = T]) == "Column", T) @@ -1777,6 +1788,13 @@ test_that("withColumn() and withColumnRenamed()", { expect_equal(length(columns(newDF)), 2) expect_equal(first(filter(newDF, df$name != "Michael"))$age, 32) + newDF <- withColumn(df, "age", 18) + expect_equal(length(columns(newDF)), 2) + expect_equal(first(newDF)$age, 18) + + expect_error(withColumn(df, "age", list("a")), + "Literal value must be atomic in length of 1") + newDF2 <- withColumnRenamed(df, "age", "newerAge") expect_equal(length(columns(newDF2)), 2) expect_equal(columns(newDF2)[1], "newerAge") --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org
