Scala Spark filter inside map

Question

I want to efficiently filter an RDD while mapping it. Is that possible?

Here is pseudocode for what I want to do:

for element in rdd:
    val opt = f(element)
    if (opt.nonEmpty) add_pair(opt.get, element)

Here is a hacky way to implement the pseudocode in Scala Spark:

rdd.map(element => (
    f(element).getOrElse(99),
    element
)).filter(tuple => tuple._1 != 99)

I wasn't able to find clean syntax to do this, so I first mapped all elements, and then filtered out the ones I don't want. Note that potentially expensive call f(element) is only computed once. If I were to filter elements before mapping (which would look more clean), then I would end up calling f twice, which is inefficient.

Please do not flag this as a duplicate. While there are similar questions, none of them actually answer this question. For example, this potential duplicate would call f twice, which is inefficient, and therefore does not answer this question.

Answer 1

You can just use flatMap:

//let's say your f returns Some(x*2) for even number and None for odd
def f(n: Int): Option[Int] = if (n % 2) Some(n*2) else None 

val rdd = sc.parallelize(List(1,2,3,4))
rdd.flatMap(f) // 4,8

// rdd.flatMap(f) or rdd.flatMap(f(_)) or rdd.flatMap(e => f(e))

And if you need to pass tuple further and filter, then just use nested map:

rdd.flatMap(e => f(e).map((_,e))) //(4,2),(8,4)