Search code examples
scalaakkaakka-stream

Flow which transforms List[A] into List[B] with flattening, given another Flow which transforms A to List[B]


I have an "inner" flow which transforms A into List[B]. I'd like to create a Flow which transforms List[A] into List[B] by invoking the "inner" flow on every element of each List[A] and then flattens the result.

For illustration, please see the test case below (the dependencies are the latest scalatest and akka-stream, A here is String, B here is Char, the "inner" flow is stringToCharacters).

The test passes, but the implementation is not idiomatic Akka Streams code since it materializes / runs the sub-streams.

Please suggest a better, idiomatic implementation which does not involve running the sub-streams.

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.ActorMaterializer
import akka.stream.scaladsl.{Flow, Sink, Source}
import org.scalatest.FlatSpec
import org.scalatest.Matchers._
import org.scalatest.concurrent.ScalaFutures.whenReady
import scala.concurrent.ExecutionContext.Implicits.global

class TestSpec extends FlatSpec {
  implicit val system = ActorSystem()
  implicit val materializer = ActorMaterializer()

  /* This is just for illustration and testing purposes, `flattenFlows` should support various inner flows */
  val stringToCharacters: Flow[String, List[Char], NotUsed] =
    Flow[String].map(x => x.toList)

  /* FIXME: I'm looking for a better implementation of the following function, please don't change the signature */
  def flattenFlows[A, B](innerFlow: Flow[A, List[B], NotUsed]): Flow[List[A], List[B], NotUsed] = {
    Flow[List[A]]
      .mapAsync(Int.MaxValue) { x =>
        Source(x).via(innerFlow).runWith(Sink.seq).map(_.flatten.toList)
      }
  }

  it should "flatten flows" in {
    val input = List(
      List("ab", "cd"),
      List("ef", "gh")
    )

    val result = Source(input).via(flattenFlows(stringToCharacters)).runWith(Sink.seq)

    val expected = List(
      List('a', 'b', 'c', 'd'),
      List('e', 'f', 'g', 'h')
    )

    whenReady(result) { x =>
      x.toList shouldEqual expected
    }
  }
}

Solution

  • You could use a combination of flatMapConcat and fold

      def flattenFlows[A, B](innerFlow: Flow[A, List[B], NotUsed]): Flow[List[A], List[B], NotUsed] =
        Flow[List[A]].flatMapConcat{ x => Source(x).via(innerFlow).fold(List.empty[B])(_ ::: _)}