fs2 `groupBy/partitions`

groupBy.scala

 // Grows with the number of distinct `K`
  def partitions[F[_], A, K](selector: A => F[K])(implicit F: Effect[F], ec: ExecutionContext)  : Pipe[F, A, (K, Stream[F, A])] = in => 
  Stream.eval(async.refOf[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]]
            _ <- st.modify(_ + (key -> newQ))
                _ <- newQ.enqueue1(el.some)
              } yield (key -> newQ.dequeue.unNoneTerminate).some
            )(_.enqueue1(el.some) as None)
        }.flatten
      }.unNone.onFinalize(cleanup)
    }

  import ExecutionContext.Implicits.global

  def selector(i: Int): IO[Int] =
    IO.pure(i % 3)

  def flakiness[A]: Pipe[IO, A, A] = in => {
    def wait = IO(scala.util.Random.nextInt(500)).flatMap(d => Timer[IO].sleep(d.millis))

    Stream.repeatEval(wait).zip(in).map(_._2)
  }

  def example =
    Stream.range(1, 100).covary[IO].through(partitions(selector)).map {
      case (k, st) => st.tupleLeft(k).through(flakiness).through(Sink.showLinesStdOut)
    }.joinUnbounded
      .compile.drain.unsafeRunSync

FS2 1.0.0 updated version: https://gist.github.com/kiambogo/8247a7bbf79f00414d1489b7e6fc90d0