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

maropu · Jun 22, 2015 · c196173 · c196173
1 parent b1f3a48
commit c196173
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/R/pkg/R/column.R b/R/pkg/R/column.R
@@ -91,6 +91,16 @@ createOperator <- function(op) {
             })
 }
 
+createInOperator <- function(op) {
+  setMethod("%in%",
+            signature(x = "Column"),
+            function(x, table) {
+              table <- listToSeq(as.list(table))
+              bar <- callJMethod(x@jc, "in", table)
+              return(column(bar))
+            })
+}
+
 createColumnFunction1 <- function(name) {
   setMethod(name,
             signature(x = "Column"),
@@ -139,6 +149,7 @@ createBinaryMathfunctions <- function(name) {
 }
 
 createMethods <- function() {
+  createInOperator()
   for (op in names(operators)) {
     createOperator(op)
   }

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", {