[SPARK-8431] [SPARKR] Add in operator to DataFrame Column in SparkR

[[SPARK-8431] Add in operator to DataFrame Column in SparkR - ASF JIRA](https://issues.apache.org/jira/browse/SPARK-8431) Author: Yu ISHIKAWA <[email protected]> Closes apache#6941 from yu-iskw/SPARK-8431 and squashes the following commits: 1f64423 [Yu ISHIKAWA] Modify the comment f4309a7 [Yu ISHIKAWA] Make a `setMethod` for `%in%` be independent 6e37936 [Yu ISHIKAWA] Modify a variable name c196173 [Yu ISHIKAWA] [SPARK-8431][SparkR] Add in operator to DataFrame Column in SparkR
mrares · Jun 23, 2015 · d4f6335 · d4f6335
1 parent 164fe2a
commit d4f6335
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/R/pkg/R/column.R b/R/pkg/R/column.R
@@ -210,6 +210,22 @@ setMethod("cast",
             }
           })
 
+#' Match a column with given values.
+#'
+#' @rdname column
+#' @return a matched values as a result of comparing with given values.
+#' \dontrun{
+#'   filter(df, "age in (10, 30)")
+#'   where(df, df$age %in% c(10, 30))
+#' }
+setMethod("%in%",
+          signature(x = "Column"),
+          function(x, table) {
+            table <- listToSeq(as.list(table))
+            jc <- callJMethod(x@jc, "in", table)
+            return(column(jc))
+          })
+
 #' Approx Count Distinct
 #'
 #' @rdname column

diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
@@ -693,6 +693,16 @@ test_that("filter() on a DataFrame", {
   filtered2 <- where(df, df$name != "Michael")
   expect_true(count(filtered2) == 2)
   expect_true(collect(filtered2)$age[2] == 19)
+
+  # test suites for %in%
+  filtered3 <- filter(df, "age in (19)")
+  expect_equal(count(filtered3), 1)
+  filtered4 <- filter(df, "age in (19, 30)")
+  expect_equal(count(filtered4), 2)
+  filtered5 <- where(df, df$age %in% c(19))
+  expect_equal(count(filtered5), 1)
+  filtered6 <- where(df, df$age %in% c(19, 30))
+  expect_equal(count(filtered6), 2)
 })
 
 test_that("join() on a DataFrame", {