Merge pull request apache#115 from sun-rui/SPARKR-130

[SPARKR-130] Add persist(storageLevel) API to RDD.

Author: shivaram

pkg/NAMESPACE

@@ -34,6 +34,7 @@ exportMethods(
               "minimum",
               "numPartitions",
               "partitionBy",
+              "persist",
               "reduce",
               "reduceByKey",
               "sampleRDD",

pkg/R/RDD.R

@@ -173,6 +173,57 @@ setMethod("cache",
             rdd
           })
 
+#' Persist an RDD
+#'
+#' Persist this RDD with the specified storage level. For details of the
+#' supported storage levels, refer to
+#' http://spark.apache.org/docs/latest/programming-guide.html#rdd-persistence.
+#'
+#' @param rdd The RDD to persist
+#' @param newLevel The new storage level to be assigned
+#' @rdname persist
+#' @export
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' persist(rdd, "MEMORY_AND_DISK")
+#'}
+setGeneric("persist", function(rdd, newLevel) { standardGeneric("persist") })
+
+#' @rdname persist
+#' @aliases persist,RDD-method
+setMethod("persist",
+          signature(rdd = "RDD", newLevel = "character"),
+          function(rdd, newLevel = c("DISK_ONLY",
+                                     "DISK_ONLY_2",
+                                     "MEMORY_AND_DISK",
+                                     "MEMORY_AND_DISK_2",
+                                     "MEMORY_AND_DISK_SER",
+                                     "MEMORY_AND_DISK_SER_2",
+                                     "MEMORY_ONLY",
+                                     "MEMORY_ONLY_2",
+                                     "MEMORY_ONLY_SER",
+                                     "MEMORY_ONLY_SER_2",
+                                     "OFF_HEAP")) {
+            match.arg(newLevel)
+            storageLevel <- switch(newLevel,
+              "DISK_ONLY" = J("org.apache.spark.storage.StorageLevel")$DISK_ONLY(),
+              "DISK_ONLY_2" = J("org.apache.spark.storage.StorageLevel")$DISK_ONLY_2(),
+              "MEMORY_AND_DISK" = J("org.apache.spark.storage.StorageLevel")$MEMORY_AND_DISK(),
+              "MEMORY_AND_DISK_2" = J("org.apache.spark.storage.StorageLevel")$MEMORY_AND_DISK_2(),
+              "MEMORY_AND_DISK_SER" = J("org.apache.spark.storage.StorageLevel")$MEMORY_AND_DISK_SER(),
+              "MEMORY_AND_DISK_SER_2" = J("org.apache.spark.storage.StorageLevel")$MEMORY_AND_DISK_SER_2(),
+              "MEMORY_ONLY" = J("org.apache.spark.storage.StorageLevel")$MEMORY_ONLY(),
+              "MEMORY_ONLY_2" = J("org.apache.spark.storage.StorageLevel")$MEMORY_ONLY_2(),
+              "MEMORY_ONLY_SER" = J("org.apache.spark.storage.StorageLevel")$MEMORY_ONLY_SER(),
+              "MEMORY_ONLY_SER_2" = J("org.apache.spark.storage.StorageLevel")$MEMORY_ONLY_SER_2(),
+              "OFF_HEAP" = J("org.apache.spark.storage.StorageLevel")$OFF_HEAP())
+            
+            .jcall(getJRDD(rdd), "Lorg/apache/spark/api/java/JavaRDD;", "persist", storageLevel)
+            rdd@env$isCached <- TRUE
+            rdd
+          })
 
 #' Unpersist an RDD
 #'

pkg/inst/tests/test_rdd.R

@@ -89,7 +89,7 @@ test_that("several transformations on RDD (a benchmark on PipelinedRDD)", {
   collect(rdd2)
 })
 
-test_that("PipelinedRDD support actions: cache(), unpersist(), checkpoint()", {
+test_that("PipelinedRDD support actions: cache(), persist(), unpersist(), checkpoint()", {
   # RDD
   rdd2 <- rdd
   # PipelinedRDD
@@ -107,6 +107,14 @@ test_that("PipelinedRDD support actions: cache(), unpersist(), checkpoint()", {
   unpersist(rdd2)
   expect_false(rdd2@env$isCached)
 
+  persist(rdd2, "MEMORY_AND_DISK")
+  expect_true(rdd2@env$isCached)
+  rdd2 <- lapply(rdd2, function(x) x)
+  expect_false(rdd2@env$isCached)
+
+  unpersist(rdd2)
+  expect_false(rdd2@env$isCached)
+
   setCheckpointDir(sc, "checkpoints")
   checkpoint(rdd2)
   expect_true(rdd2@env$isCheckpointed)

pkg/man/persist.Rd

@@ -0,0 +1,33 @@
+% Generated by roxygen2 (4.0.2): do not edit by hand
+\docType{methods}
+\name{persist}
+\alias{persist}
+\alias{persist,RDD,character-method}
+\alias{persist,RDD-method}
+\title{Persist an RDD}
+\usage{
+persist(rdd, newLevel)
+
+\S4method{persist}{RDD,character}(rdd, newLevel = c("DISK_ONLY",
+  "DISK_ONLY_2", "MEMORY_AND_DISK", "MEMORY_AND_DISK_2", "MEMORY_AND_DISK_SER",
+  "MEMORY_AND_DISK_SER_2", "MEMORY_ONLY", "MEMORY_ONLY_2", "MEMORY_ONLY_SER",
+  "MEMORY_ONLY_SER_2", "OFF_HEAP"))
+}
+\arguments{
+\item{rdd}{The RDD to persist}
+
+\item{newLevel}{The new storage level to be assigned}
+}
+\description{
+Persist this RDD with the specified storage level. For details of the
+supported storage levels, refer to
+http://spark.apache.org/docs/latest/programming-guide.html#rdd-persistence.
+}
+\examples{
+\dontrun{
+sc <- sparkR.init()
+rdd <- parallelize(sc, 1:10, 2L)
+persist(rdd, "MEMORY_AND_DISK")
+}
+}
+