Akka Streams custom merge - scala

Question

I'm new to akka-streams and not sure how to approach this problem.
I have 3 source streams are sorted by a sequence ID. I want to group the values together which have the same ID. Values in each stream may be missing or duplicated. If one stream is a faster producer than the rest, it should get backpressured.
case class A(id: Int)
case class B(id: Int)
case class C(id: Int)
case class Merged(as: List[A], bs: List[B], cs: List[C])
import akka.stream._
import akka.stream.scaladsl._
val as = Source(List(A(1), A(2), A(3), A(4), A(5)))
val bs = Source(List(B(1), B(2), B(3), B(4), B(5)))
val cs = Source(List(C(1), C(1), C(3), C(4)))
val merged = ???
// value 1: Merged(List(A(1)), List(B(1)), List(C(1), C(1)))
// value 2: Merged(List(A(2)), List(B(2)), Nil)
// value 3: Merged(List(A(3)), List(B(3)), List(C(3)))
// value 4: Merged(List(A(4)), List(B(4)), List(C(4)))
// value 5: Merged(List(A(5)), List(B(5)), Nil)
// (end of stream)

Answer 1

this question is old but I was trying to find a solution for that and I only encountered the rocks to the path at the lightbend forum, but not a working use case. So I decided to implement and post here my example.
I created 3 sources sourceA, sourceB, and sourceC which emit events in a different speed using .throttle(). Then I created a RunnableGraph where I merge the sources using Merge and I the output to my WindowGroupEventFlow Flow that I implemented based on a sliding window of number of events. This is the graph:
sourceA ~> mergeShape.in(0)
sourceB ~> mergeShape.in(1)
sourceC ~> mergeShape.in(2)
mergeShape.out ~> windowFlowShape ~> sinkShape
The classes that I am using on the sources are these:
object Domain {
sealed abstract class Z(val id: Int, val value: String)
case class A(override val id: Int, override val value: String = "A") extends Z(id, value)
case class B(override val id: Int, override val value: String = "B") extends Z(id, value)
case class C(override val id: Int, override val value: String = "C") extends Z(id, value)
case class ABC(override val id: Int, override val value: String) extends Z(id, value)
}
and this is the WindowGroupEventFlow Flow that I created to group the events:
// step 0: define the shape
class WindowGroupEventFlow(maxBatchSize: Int) extends GraphStage[FlowShape[Domain.Z, Domain.Z]] {
// step 1: define the ports and the component-specific members
val in = Inlet[Domain.Z]("WindowGroupEventFlow.in")
val out = Outlet[Domain.Z]("WindowGroupEventFlow.out")
// step 3: create the logic
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
// mutable state
val batch = new mutable.Queue[Domain.Z]
var count = 0
// var result = ""
// step 4: define mutable state implement my logic here
setHandler(in, new InHandler {
override def onPush(): Unit = {
try {
val nextElement = grab(in)
batch.enqueue(nextElement)
count += 1
// If window finished we have to dequeue all elements
if (count >= maxBatchSize) {
println("************ window finished - dequeuing elements ************")
var result = Map[Int, Domain.Z]()
val list = batch.dequeueAll(_ => true).to[collection.immutable.Iterable]
list.foreach { tuple =>
if (result.contains(tuple.id)) {
val abc = result.get(tuple.id)
val value = abc.get.value + tuple.value
val z: Domain.Z = Domain.ABC(tuple.id, value)
result += (tuple.id -> z)
} else {
val z: Domain.Z = Domain.ABC(tuple.id, tuple.value)
result += (tuple.id -> z)
}
}
val finalResult: collection.immutable.Iterable[Domain.Z] = result.map(p => p._2)
emitMultiple(out, finalResult)
count = 0
} else {
pull(in) // send demand upstream signal, asking for another element
}
} catch {
case e: Throwable => failStage(e)
}
}
})
setHandler(out, new OutHandler {
override def onPull(): Unit = {
pull(in)
}
})
}
// step 2: construct a new shape
override def shape: FlowShape[Domain.Z, Domain.Z] = FlowShape[Domain.Z, Domain.Z](in, out)
}
and this is how I run everything:
object WindowGroupEventFlow {
def main(args: Array[String]): Unit = {
run()
}
def run() = {
implicit val system = ActorSystem("WindowGroupEventFlow")
import Domain._
val sourceA = Source(List(A(1), A(2), A(3), A(1), A(2), A(3), A(1), A(2), A(3), A(1))).throttle(3, 1 second)
val sourceB = Source(List(B(1), B(2), B(1), B(2), B(1), B(2), B(1), B(2), B(1), B(2))).throttle(2, 1 second)
val sourceC = Source(List(C(1), C(2), C(3), C(4))).throttle(1, 1 second)
// Step 1 - setting up the fundamental for a stream graph
val windowRunnableGraph = RunnableGraph.fromGraph(
GraphDSL.create() { implicit builder =>
import GraphDSL.Implicits._
// Step 2 - create shapes
val mergeShape = builder.add(Merge[Domain.Z](3))
val windowEventFlow = Flow.fromGraph(new WindowGroupEventFlow(5))
val windowFlowShape = builder.add(windowEventFlow)
val sinkShape = builder.add(Sink.foreach[Domain.Z](x => println(s"sink: $x")))
// Step 3 - tying up the components
sourceA ~> mergeShape.in(0)
sourceB ~> mergeShape.in(1)
sourceC ~> mergeShape.in(2)
mergeShape.out ~> windowFlowShape ~> sinkShape
// Step 4 - return the shape
ClosedShape
}
)
// run the graph and materialize it
val graph = windowRunnableGraph.run()
}
}
you can see on the output how I am grouping the elements with same ID:
sink: ABC(1,ABC)
sink: ABC(2,AB)
************ window finished - dequeuing elements ************
sink: ABC(3,A)
sink: ABC(1,BA)
sink: ABC(2,CA)
************ window finished - dequeuing elements ************
sink: ABC(2,B)
sink: ABC(3,AC)
sink: ABC(1,BA)
************ window finished - dequeuing elements ************
sink: ABC(2,AB)
sink: ABC(3,A)
sink: ABC(1,BA)