Merge pull request apache#118 from sun-rui/saveAsTextFile

[SPARKR-144] Implement saveAsTextFile() in the RDD class.

Author: shivaram

pkg/NAMESPACE

@@ -42,6 +42,7 @@ exportMethods(
               "reduceByKey",
               "rightOuterJoin",
               "sampleRDD",
+              "saveAsTextFile",
               "saveAsObjectFile",
               "take",
               "takeSample",

pkg/R/RDD.R

@@ -85,10 +85,10 @@ PipelinedRDD <- function(prev, func) {
 
 
 # The jrdd accessor function.
-setGeneric("getJRDD", function(rdd) { standardGeneric("getJRDD") })
+setGeneric("getJRDD", function(rdd, ...) { standardGeneric("getJRDD") })
 setMethod("getJRDD", signature(rdd = "RDD"), function(rdd) rdd@jrdd )
 setMethod("getJRDD", signature(rdd = "PipelinedRDD"),
-          function(rdd) {
+          function(rdd, dataSerialization = TRUE) {
             if (!is.null(rdd@env$jrdd_val)) {
               return(rdd@env$jrdd_val)
             }
@@ -116,17 +116,29 @@ setMethod("getJRDD", signature(rdd = "PipelinedRDD"),
 
             prev_jrdd <- rdd@prev_jrdd
 
-            rddRef <- new(J("edu.berkeley.cs.amplab.sparkr.RRDD"),
-                           prev_jrdd$rdd(),
-                           serializedFuncArr,
-                           rdd@env$serialized,
-                           depsBinArr,
-                           packageNamesArr,
-                           as.character(.sparkREnv[["libname"]]),
-                           broadcastArr,
-                           prev_jrdd$classTag())
-            # The RDD is serialized after we create a RRDD
-            rdd@env$serialized <- TRUE
+            if (dataSerialization) {
+              rddRef <- new(J("edu.berkeley.cs.amplab.sparkr.RRDD"),
+                             prev_jrdd$rdd(),
+                             serializedFuncArr,
+                             rdd@env$serialized,
+                             depsBinArr,
+                             packageNamesArr,
+                             as.character(.sparkREnv[["libname"]]),
+                             broadcastArr,
+                             prev_jrdd$classTag())
+            } else {
+              rddRef <- new(J("edu.berkeley.cs.amplab.sparkr.StringRRDD"),
+                             prev_jrdd$rdd(),
+                             serializedFuncArr,
+                             rdd@env$serialized,
+                             depsBinArr,
+                             packageNamesArr,
+                             as.character(.sparkREnv[["libname"]]),
+                             broadcastArr,
+                             prev_jrdd$classTag())
+            }
+            # Save the serialization flag after we create a RRDD
+            rdd@env$serialized <- dataSerialization
             rdd@env$jrdd_val <- rddRef$asJavaRDD()
             rdd@env$jrdd_val
           })
@@ -1091,6 +1103,34 @@ setMethod("saveAsObjectFile",
             invisible(NULL)
           })
 
+#' Save this RDD as a text file, using string representations of elements.
+#'
+#' @param rdd The RDD to save
+#' @param path The directory where the splits of the text file are saved
+#' @rdname saveAsTextFile
+#' @export
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:3)
+#' saveAsTextFile(rdd, "/tmp/sparkR-tmp")
+#'}
+setGeneric("saveAsTextFile", function(rdd, path) { standardGeneric("saveAsTextFile") })
+
+#' @rdname saveAsTextFile
+#' @aliases saveAsTextFile,RDD
+setMethod("saveAsTextFile",
+          signature(rdd = "RDD", path = "character"),
+          function(rdd, path) {
+            func <- function(x) {
+              toString(x)
+            }
+            stringRdd <- lapply(rdd, func)
+            .jcall(getJRDD(stringRdd, dataSerialization = FALSE), "V", "saveAsTextFile", path)
+            # Return nothing
+            invisible(NULL)
+          })
+
 #' Return an RDD with the keys of each tuple.
 #'
 #' @param rdd The RDD from which the keys of each tuple is returned.

pkg/inst/tests/test_textFile.R

@@ -58,3 +58,52 @@ test_that("several transformations on RDD created by textFile()", {
 
   unlink(fileName)
 })
+
+test_that("textFile() followed by a saveAsTextFile() returns the same content", {
+  fileName1 <- tempfile(pattern="spark-test", fileext=".tmp")
+  fileName2 <- tempfile(pattern="spark-test", fileext=".tmp")
+  writeLines(mockFile, fileName1)
+
+  rdd <- textFile(sc, fileName1)
+  saveAsTextFile(rdd, fileName2)
+  rdd <- textFile(sc, fileName2)
+  expect_equal(collect(rdd), as.list(mockFile))
+
+  unlink(fileName1)
+  unlink(fileName2)
+})
+
+test_that("saveAsTextFile() on a parallelized list works as expected", {
+  fileName <- tempfile(pattern="spark-test", fileext=".tmp")
+  l <- list(1, 2, 3)
+  rdd <- parallelize(sc, l)
+  saveAsTextFile(rdd, fileName)
+  rdd <- textFile(sc, fileName)
+  expect_equal(collect(rdd), lapply(l, function(x) {toString(x)}))
+
+  unlink(fileName)
+})
+
+test_that("textFile() and saveAsTextFile() word count works as expected", {
+  fileName1 <- tempfile(pattern="spark-test", fileext=".tmp")
+  fileName2 <- tempfile(pattern="spark-test", fileext=".tmp")
+  writeLines(mockFile, fileName1)
+
+  rdd <- textFile(sc, fileName1)
+
+  words <- flatMap(rdd, function(line) { strsplit(line, " ")[[1]] })
+  wordCount <- lapply(words, function(word) { list(word, 1L) })
+
+  counts <- reduceByKey(wordCount, "+", 2L)
+
+  saveAsTextFile(counts, fileName2)
+  rdd <- textFile(sc, fileName2)
+   
+  output <- collect(rdd)
+  expected <- list(list("awesome.", 1), list("Spark", 2),
+                   list("pretty.", 1), list("is", 2))
+  expect_equal(output, lapply(expected, function(x) {toString(x)}))
+  
+  unlink(fileName1)
+  unlink(fileName2)
+})

pkg/inst/worker/serialize.R

@@ -46,6 +46,10 @@ writeInt <- function(con, value) {
   writeBin(as.integer(value), con, endian="big")
 }
 
+writeStrings <- function(con, stringList) {
+  writeLines(unlist(stringList), con)
+}
+
 writeRaw <- function(con, batch, serialized = FALSE) {
   if (serialized) {
     outputSer <- batch

pkg/inst/worker/worker.R

@@ -33,8 +33,11 @@ splitIndex <- readInt(inputCon)
 execLen <- readInt(inputCon)
 execFunctionName <- unserialize(readRawLen(inputCon, execLen))
 
-# read the isSerialized bit flag
-isSerialized <- readInt(inputCon)
+# read the isInputSerialized bit flag
+isInputSerialized <- readInt(inputCon)
+
+# read the isOutputSerialized bit flag
+isOutputSerialized <- readInt(inputCon)
 
 # Redirect stdout to stderr to prevent print statements from
 # interfering with outputStream
@@ -82,16 +85,20 @@ isEmpty <- readInt(inputCon)
 if (isEmpty != 0) {
 
   if (numPartitions == -1) {
-    if (isSerialized) {
+    if (isInputSerialized) {
       # Now read as many characters as described in funcLen
       data <- readDeserialize(inputCon)
     } else {
       data <- readLines(inputCon)
     }
     output <- do.call(execFunctionName, list(splitIndex, data))
-    writeRaw(outputCon, output)
+    if (isOutputSerialized) {
+      writeRaw(outputCon, output)
+    } else {
+      writeStrings(outputCon, output)
+    }
   } else {
-    if (isSerialized) {
+    if (isInputSerialized) {
       # Now read as many characters as described in funcLen
       data <- readDeserialize(inputCon)
     } else {
@@ -127,7 +134,9 @@ if (isEmpty != 0) {
 }
 
 # End of output
-writeInt(outputCon, 0L)
+if (isOutputSerialized) {
+  writeInt(outputCon, 0L)
+}
 
 close(outputCon)