Parallel Map in Kotlin

Ever wonder how to run map in parallel using coroutines? This is how you do it.

import kotlinx.coroutines.async import kotlinx.coroutines.runBlocking import kotlinx.coroutines.coroutineScope import kotlinx.coroutines.awaitAll //sampleStart suspend fun <A, B> Iterable<A>.pmap(f: suspend (A) -> B): List<B> = coroutineScope { map { async { f(it) } }.awaitAll() } //sampleEnd

Confused? Let’s unpack it.

First we have the function signature which is pretty similar to the actual map extension function signature on Iterable. The only thing we added were the suspend keywords. One for our function and another one on the parameter.

Then we have the coroutineScope that marks the scope in which the async calls are going to be executed. This way if something goes wrong and an exception is thrown, all coroutines that were launched in this scope are cancelled. That’s the main benefit of structured concurrency.

Finally we have the actual execution which is divided in 2 steps: The first step launches a new coroutine for each function application using async. This effectively wraps the type of each element with Deferred.

In the second step we wait for all function invocations to complete and unwrap the result using awaitAll(). This is similar to doing .map { it.await() } but better because awaitAll() will fail as soon as one of the invocations fails, instead of having to sequentially wait for the call to await() on the failing deferred. For example, let’s say we call pmap with 3 elements. f(0) will complete in 2 minutes, f(1) completes in 5 minutes and f(3) fails immediately. With .map { it.await() } we’d have to wait for f(1) and f(2) completion before seeing the f(3) failure, whereas awaitAll() will know that something failed right away.

How to use it

Easy! Just like you use map:

import kotlinx.coroutines.async import kotlinx.coroutines.coroutineScope import kotlinx.coroutines.runBlocking import kotlinx.coroutines.awaitAll suspend fun <A, B> Iterable<A>.pmap(f: suspend (A) -> B): List<B> = coroutineScope { map { async { f(it) } }.awaitAll() } //sampleStart fun main(args: Array<String>) = runBlocking { println((1..100).pmap { it * 2 }) } //sampleEnd

(Psst! I’m using Kotlin Playground so you can actually run this code!)

Prove that it’s running in parallel

Ok so let’s resort to the good old delay to prove that this is actually running in parallel. We are going to add a delay of 1 second on each multiplication and measure the time it takes to run.

Running over 100 elements the result should be: close to 1,000 milliseconds if it’s running in parallel and close to 100,000 milliseconds if it’s running sequentially.

import kotlinx.coroutines.async import kotlinx.coroutines.runBlocking import kotlin.system.measureTimeMillis import kotlinx.coroutines.delay import kotlinx.coroutines.coroutineScope import kotlinx.coroutines.awaitAll suspend fun <A, B> Iterable<A>.pmap(f: suspend (A) -> B): List<B> = coroutineScope { map { async { f(it) } }.awaitAll() } //sampleStart fun main(args: Array<String>) = runBlocking { val time = measureTimeMillis { val output = (1..100).pmap { delay(1000) it * 2 } println(output) } println("Total time: $time") } //sampleEnd

Why not runBlocking?

A previous iteration of this article proposed the following solution:

import kotlinx.coroutines.async import kotlinx.coroutines.runBlocking import kotlinx.coroutines.awaitAll //sampleStart fun <A, B> Iterable<A>.pmapOld(f: suspend (A) -> B): List<B> = runBlocking { map { async { f(it) } }.awaitAll() } //sampleEnd

And the description said:

Then we have the runBlocking which lets us bridge the blocking code with the coroutine world. As the name suggests this will block the current thread until everything inside the block finishes executing. Which is exactly what we want.

Fortunately Gildor commented on why this a bad idea. First, we were not using Structured Concurrency, so an invocation of f could fail and all other executions would continue unfazed. And also we were not playing nice with the rest of the code. By using runBlocking we were forcefully blocking the thread until the whole execution of pmap finishes, instead of letting the caller decide how the execution should go.

---

Comments