Feature request: groupWeighted

[fugafree]

This is a question about an idea: introducing Stream.groupWeighted. It would output Chunks limited by the weight of the elements. Weight would be provided by a weight function.
Similar to Akka's groupedWeightedWithin, but without the time config.
Or even a Stream.groupWeightedWithin method, similar to the already existing Stream.groupWithin.
What do you think?

[fugafree]

Or it could be named chunkWeighted, I don't know which one would be better.

[RobertoUa]

something like this?

  implicit class StreamOps[F[_], O](private val s: Stream[F, O]) extends AnyVal {

    def groupWeighedWithin(timeout: FiniteDuration, limit: Long, weight: O => Long)(implicit
      F: Temporal[F]
    ): Stream[F, Chunk[O]] =
      s.pull
        .timed { timedPull =>
          def go(timedPull: Pull.Timed[F, O], buffer: Chunk[O], bufferSize: Long): Pull[F, O, Unit] =
            timedPull.timeout(timeout) >>
              timedPull.uncons.flatMap {
                case Some((Right(elems), next)) =>
                  val size     = elems.map(weight).sumAll
                  val nextSize = bufferSize + size
                  if (nextSize >= limit)
                    Pull.output(buffer) >> go(next, elems, size)
                  else
                    go(next, buffer ++ elems, nextSize)
                case Some((Left(_), next)) =>
                  Pull.output(buffer) >> go(next, Chunk.empty, 0)
                case None =>
                  Pull.output(buffer) >> Pull.done
              }

          go(timedPull, Chunk.empty, 0)
        }
        .stream
        .chunks
  }

[fugafree]

Wow, looks promising, thank you very much. I'm going to dig more into this, as it contains a couple of new concepts for me, Pulls are a bit new to me.

[fugafree]

For my use-case, no need for timed pull, a groupWeighted would be a simpler and sufficient first step.

Observations:

• I think uncons1 should be used instead of uncons, because it could return a chunk, which has already a greater sum weight, than allowed. And we can not even split the chunk, like in unconsN, becuase there is no splitAtWeight or similar.
• also even a single element can be larger than limit, that can be handled easily inside go

My proposal:

implicit class StreamOps[F[_]: RaiseThrowable, O](private val s: Stream[F, O]) extends AnyVal {

  def groupWeighed(
      limit: Long,
      weight: O => Long
  ): Stream[F, Chunk[O]] = {
    def go(stream: Stream[F, O], buffer: Chunk[O], bufferSize: Long): Pull[F, O, Unit] =
      stream.pull.uncons1.flatMap {
        case Some((elem, next)) =>
          val size     = weight(elem)
          val nextSize = bufferSize + size
          if (size >= limit)
            Pull.raiseError[F](new RuntimeException("Too large element"))
          else if (nextSize >= limit)
            Pull.output(buffer) >> go(next, Chunk(elem), size)
          else
            go(next, buffer ++ Chunk(elem), nextSize)
        case None =>
          Pull.output(buffer) >> Pull.done
      }

    go(s, Chunk.empty, 0).stream.chunks
  }

}

How does this look like? Since this is my first pull transformation I can not tell if this is how it should look like or not. Another difference from your code is that it uses a Stream as the first param of go, as uncons1 returns a Stream instead of a Pull. So either here or in the go we will need to call .pull.uncons1 always on the returned Stream. I don't know if that's okey or not.