Merge pull request apache#153 from hqzizania/master

[SPARKR-160] Support collectAsMap()

Author: shivaram

pkg/NAMESPACE

@@ -6,6 +6,7 @@ exportMethods(
               "checkpoint",
               "cogroup",
               "collect",
+              "collectAsMap",
               "collectPartition",
               "combineByKey",
               "count",

pkg/R/RDD.R

@@ -358,6 +358,7 @@ setMethod("collect",
             convertJListToRList(collected, flatten)
           })
 
+
 #' @rdname collect-methods
 #' @export
 #' @description
@@ -382,6 +383,29 @@ setMethod("collectPartition",
             convertJListToRList(jList, flatten = TRUE)
           })
 
+#' @rdname collect-methods
+#' @export
+#' @description
+#' \code{collectAsMap} returns a named list as a map that contains all of the elements
+#' in a key-value pair RDD. 
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 2), list(3, 4)), 2L)
+#' collectAsMap(rdd) # list(`1` = 2, `3` = 4)
+#'}
+setGeneric("collectAsMap", function(rdd) { standardGeneric("collectAsMap") })
+
+#' @rdname collect-methods
+#' @aliases collectAsMap,RDD-method
+setMethod("collectAsMap",
+          signature(rdd = "RDD"),
+          function(rdd) {
+            pairList <- collect(rdd)
+            map <- new.env()
+            lapply(pairList, function(x) { assign(as.character(x[[1]]), x[[2]], envir = map) })
+            as.list(map)
+          })
 
 #' Look up elements of a key in an RDD
 #'

pkg/inst/tests/test_rdd.R

@@ -373,3 +373,22 @@ test_that("fullOuterJoin() on pairwise RDDs", {
   expect_equal(sortKeyValueList(actual),
                sortKeyValueList(list(list("a", list(1, NULL)), list("b", list(2, NULL)), list("d", list(NULL, 4)), list("c", list(NULL, 3)))))
 })
+
+test_that("collectAsMap() on a pairwise RDD", {
+  rdd <- parallelize(sc, list(list(1, 2), list(3, 4)))
+  vals <- collectAsMap(rdd)
+  expect_equal(vals, list(`1` = 2, `3` = 4))
+
+  rdd <- parallelize(sc, list(list("a", 1), list("b", 2)))
+  vals <- collectAsMap(rdd)
+  expect_equal(vals, list(a = 1, b = 2))
+ 
+  rdd <- parallelize(sc, list(list(1.1, 2.2), list(1.2, 2.4)))
+  vals <- collectAsMap(rdd)
+  expect_equal(vals, list(`1.1` = 2.2, `1.2` = 2.4))
+ 
+  rdd <- parallelize(sc, list(list(1, "a"), list(2, "b")))
+  vals <- collectAsMap(rdd)
+  expect_equal(vals, list(`1` = "a", `2` = "b"))
+})
+

pkg/man/collect-methods.Rd

@@ -1,8 +1,11 @@
-% Generated by roxygen2 (4.0.2): do not edit by hand
+% Generated by roxygen2 (4.1.0): do not edit by hand
+% Please edit documentation in R/RDD.R
 \docType{methods}
 \name{collect}
 \alias{collect}
 \alias{collect,RDD-method}
+\alias{collectAsMap}
+\alias{collectAsMap,RDD-method}
 \alias{collectPartition}
 \alias{collectPartition,RDD,integer-method}
 \alias{collectPartition,integer,RDD-method}
@@ -15,6 +18,10 @@ collect(rdd, ...)
 collectPartition(rdd, partitionId)
 
 \S4method{collectPartition}{RDD,integer}(rdd, partitionId)
+
+collectAsMap(rdd)
+
+\S4method{collectAsMap}{RDD}(rdd)
 }
 \arguments{
 \item{rdd}{The RDD to collect}
@@ -33,6 +40,9 @@ a list containing elements in the RDD
 
 \code{collectPartition} returns a list that contains all of the elements
 in the specified partition of the RDD.
+
+\code{collectAsMap} returns a named list as a map that contains all of the elements
+in a key-value pair RDD.
 }
 \examples{
 \dontrun{
@@ -41,5 +51,10 @@ rdd <- parallelize(sc, 1:10, 2L)
 collect(rdd) # list from 1 to 10
 collectPartition(rdd, 0L) # list from 1 to 5
 }
+\dontrun{
+sc <- sparkR.init()
+rdd <- parallelize(sc, list(list(1, 2), list(3, 4)), 2L)
+collectAsMap(rdd) # list(`1` = 2, `3` = 4)
+}
 }