Merge pull request apache#84 from sun-rui/SPARKR-94

[SPARKR-94] Add a  method to get an element of a pair RDD object by key.

Author: shivaram

pkg/NAMESPACE

@@ -15,6 +15,7 @@ exportMethods(
               "lapply",
               "lapplyPartition",
               "lapplyPartitionsWithIndex",
+              "lookup",
               "map",
               "mapPartitions",
               "mapPartitionsWithIndex",

pkg/R/RDD.R

@@ -291,6 +291,38 @@ setMethod("collectPartition",
           })
 
 
+#' Look up elements of a key in an RDD
+#'
+#' @description
+#' \code{lookup} returns a list of values in this RDD for key key.
+#'
+#' @param rdd The RDD to collect
+#' @param key The key to look up for
+#' @return a list of values in this RDD for key key
+#' @rdname lookup
+#' @export
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(c(1, 1), c(2, 2), c(1, 3))
+#' rdd <- parallelize(sc, pairs)
+#' lookup(rdd, 1) # list(1, 3)
+#'}
+setGeneric("lookup", function(rdd, key) { standardGeneric("lookup") })
+
+#' @rdname lookup
+#' @aliases lookup,RDD-method
+setMethod("lookup",
+          signature(rdd = "RDD", key = "ANY"),
+          function(rdd, key) {
+            partitionFunc <- function(part) {
+              filtered <- part[unlist(lapply(part, function(x) { identical(key, x[[1]]) }))]
+              lapply(filtered, function(x) { x[[2]] })
+            }
+            valsRDD <- lapplyPartition(rdd, partitionFunc)
+            collect(valsRDD)
+          })
+
 #' Return the number of elements in the RDD.
 #'
 #' @param rdd The RDD to count

pkg/inst/tests/test_rdd.R

@@ -7,6 +7,9 @@ sc <- sparkR.init()
 nums <- 1:10
 rdd <- parallelize(sc, nums, 2L)
 
+intPairs <- list(list(1L, -1), list(2L, 100), list(2L, 1), list(1L, 200))
+intRdd <- parallelize(sc, intPairs, 2L)
+
 test_that("count and length on RDD", {
    expect_equal(count(rdd), 10)
    expect_equal(length(rdd), 10)
@@ -30,6 +33,14 @@ test_that("mapPartitions on RDD", {
   expect_equal(actual, list(15, 40))
 })
 
+test_that("lookup on RDD", {
+  vals <- lookup(intRdd, 1L)
+  expect_equal(vals, list(-1, 200))
+  
+  vals <- lookup(intRdd, 3L)
+  expect_equal(vals, list())
+})
+
 test_that("several transformations on RDD (a benchmark on PipelinedRDD)", {
   rdd2 <- rdd
   for (i in 1:12)

pkg/man/lookup.Rd

@@ -0,0 +1,31 @@
+% Generated by roxygen2 (4.0.2): do not edit by hand
+\docType{methods}
+\name{lookup}
+\alias{lookup}
+\alias{lookup,RDD-method}
+\title{Look up elements of a key in an RDD}
+\usage{
+lookup(rdd, key)
+
+\S4method{lookup}{RDD}(rdd, key)
+}
+\arguments{
+\item{rdd}{The RDD to collect}
+
+\item{key}{The key to look up for}
+}
+\value{
+a list of values in this RDD for key key
+}
+\description{
+\code{lookup} returns a list of values in this RDD for key key.
+}
+\examples{
+\dontrun{
+sc <- sparkR.init()
+pairs <- list(c(1, 1), c(2, 2), c(1, 3))
+rdd <- parallelize(sc, pairs)
+lookup(rdd, 1) # list(1, 3)
+}
+}
+