Skip to content

Commit

Permalink
[SPARK-10904] [SPARKR] Fix to support select(df, c("col1", "col2"))
Browse files Browse the repository at this point in the history
The fix is to coerce `c("a", "b")` into a list such that it could be serialized to call JVM with.

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #8961 from felixcheung/rselect.

(cherry picked from commit 721e8b5)
Signed-off-by: Shivaram Venkataraman <shivaram@cs.berkeley.edu>
  • Loading branch information
felixcheung authored and shivaram committed Oct 4, 2015
1 parent cbc6aec commit 8836ac3
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 3 deletions.
14 changes: 11 additions & 3 deletions R/pkg/R/DataFrame.R
Expand Up @@ -1044,12 +1044,20 @@ setMethod("subset", signature(x = "DataFrame"),
#' select(df, c("col1", "col2"))
#' select(df, list(df$name, df$age + 1))
#' # Similar to R data frames columns can also be selected using `$`
#' df$age
#' df[,df$age]
#' }
setMethod("select", signature(x = "DataFrame", col = "character"),
function(x, col, ...) {
sdf <- callJMethod(x@sdf, "select", col, toSeq(...))
dataFrame(sdf)
if (length(col) > 1) {
if (length(list(...)) > 0) {
stop("To select multiple columns, use a character vector or list for col")
}

select(x, as.list(col))
} else {
sdf <- callJMethod(x@sdf, "select", col, toSeq(...))
dataFrame(sdf)
}
})

#' @rdname select
Expand Down
7 changes: 7 additions & 0 deletions R/pkg/inst/tests/test_sparkSQL.R
Expand Up @@ -585,6 +585,13 @@ test_that("select with column", {
expect_equal(columns(df3), c("x"))
expect_equal(count(df3), 3)
expect_equal(collect(select(df3, "x"))[[1, 1]], "x")

df4 <- select(df, c("name", "age"))
expect_equal(columns(df4), c("name", "age"))
expect_equal(count(df4), 3)

expect_error(select(df, c("name", "age"), "name"),
"To select multiple columns, use a character vector or list for col")
})

test_that("subsetting", {
Expand Down

0 comments on commit 8836ac3

Please sign in to comment.