Grouping a stream of elements into multiple streams

user_8990329 0
Viewed 598

Let's say we have a case class MyCaseClass(name: String, value: Int). Given an fs2.Stream[F, MyCaseClass] I want to group elements with the same name

val sourceStream: fs2.Stream[F, MyCaseClass] = //
val groupedSameNameStream: fs2.Stream[F, fs2.Stream[F, MyCaseClass]] = //

The reason I need to do this is I want to apply effectfful transformation

val transform: MyCaseClass => F[Unit] = //

to all elements of a stream and in case one group fails the other should keep working.

Is something like this possible to do?

fs2 scala-cats scala functional-programming
1 Answers

This is possible, with caveats.

It's relatively straightforward to do this if you accept having a Map with an unbounded number of keys, and an unbounded number of associated Queues for each.

We've use code based on a gist by github user kiambogo in production (though ours has been tweaked), and it works fine:

import fs2.concurrent.Queue
import cats.implicits._
import cats.effect.Concurrent
import cats.effect.concurrent.Ref
 
def groupBy[F[_], A, K](selector: A => F[K])(implicit F: Concurrent[F]): Pipe[F, A, (K, Stream[F, A])] = {
  in =>
  Stream.eval(Ref.of[F, Map[K, Queue[F, Option[A]]]](Map.empty)).flatMap { st =>
    val cleanup = {
      import alleycats.std.all._
      st.get.flatMap(_.traverse_(_.enqueue1(None)))
    }

    (in ++ Stream.eval_(cleanup))
      .evalMap { el =>
        (selector(el), st.get).mapN { (key, queues) =>
          queues.get(key).fold {
            for {
              newQ <- Queue.unbounded[F, Option[A]] // Create a new queue
              _ <- st.modify(x => (x + (key -> newQ), x)) // Update the ref of queues
              _ <- newQ.enqueue1(el.some)
            } yield (key -> newQ.dequeue.unNoneTerminate).some
          }(_.enqueue1(el.some) as None)
        }.flatten
      }.unNone.onFinalize(cleanup)
  }
}

If we assume an overhead of 64 bytes for each Map entry (I believe this is very overestimated) then a cardinality of 100,000 unique keys gives us approximately 6.1MiB - well within reasonable size for a jvm process.

user_3503510
Related