zio · kyri-petrou · Mar 28, 2025 · Mar 14, 2025 · Mar 14, 2025 · Mar 14, 2025
diff --git a/streams-tests/shared/src/test/scala/zio/stream/ZStreamSpec.scala b/streams-tests/shared/src/test/scala/zio/stream/ZStreamSpec.scala
@@ -2735,6 +2735,39 @@ object ZStreamSpec extends ZIOBaseSpec {
             )(equalTo(Chunk(Right(1), Right(2), Left("boom"))))
           }
         ),
+        suite("mapZIOChunked")(
+          test("ZIO#foreach equivalence when error-free") {
+            check(Gen.small(Gen.listOfN(_)(Gen.byte)), Gen.function(Gen.successes(Gen.byte))) { (data, f) =>
+              val s = ZStream.fromIterable(data)
+
+              for {
+                l <- s.mapZIOChunked(f).runCollect
+                r <- ZIO.foreach(data)(f)
+              } yield assert(l.toList)(equalTo(r))
+            }
+          },
+          test("retains chunks") {
+            ZStream
+              .fromChunks(Chunk(1, 2))
+              .mapZIOChunked(ZIO.succeed(_))
+              .chunks
+              .runCount
+              .map(count => assertTrue(count == 1))
+          },
+          test("retains chunks up to error") {
+            assertZIO(
+              ZStream
+                .fromChunks(Chunk(1, 2), Chunk(3, 4, 5))
+                .mapZIOChunked {
+                  case 4 => ZIO.fail("boom")
+                  case x => ZIO.succeed(x)
+                }
+                .either
+                .chunks
+                .runCollect
+            )(equalTo(Chunk(Chunk(Right(1), Right(2)), Chunk(Right(3)), Chunk(Left("boom")))))
+          }
+        ),
         suite("mapZIOPar")(
           test("foreachParN equivalence") {
             checkN(10)(Gen.small(Gen.listOfN(_)(Gen.byte)), Gen.function(Gen.successes(Gen.byte))) { (data, f) =>

diff --git a/streams/shared/src/main/scala/zio/stream/ZPipeline.scala b/streams/shared/src/main/scala/zio/stream/ZPipeline.scala
@@ -482,6 +482,25 @@ final class ZPipeline[-Env, +Err, -In, +Out] private (
   )(implicit trace: Trace): ZPipeline[Env2, Err2, In, Out2] =
     self >>> ZPipeline.mapZIO(f)
 
+  /**
+   * Creates a pipeline that maps over elements of the stream with the specified
+   * effectful function.
+   *
+   * Unlike `mapZIO` processing is done chunk by chunk. This means that
+   * `mapZIOChunked` provides weaker guarantees than `mapZIO`. While
+   * `stream.mapZIO(f).mapZIO(g)` is guaranteed to be equivalent to
+   * `stream.mapZIO(x => f(x).flatMap(g))`, the same is not true for
+   * `mapZIOChunked`. For example, `mapZIO` guarantees that the first element of
+   * a stream will first be processed with `f` and then `g` before the second
+   * element is processed with `f`. `mapZIOChunked` may process the first two
+   * elements with `f` and only then move on to process the first element with
+   * `g`.
+   */
+  def mapZIOChunked[Env2 <: Env, Err2 >: Err, Out2](
+    f: Out => ZIO[Env2, Err2, Out2]
+  )(implicit trace: Trace): ZPipeline[Env2, Err2, In, Out2] =
+    self >>> ZPipeline.mapZIOChunked(f)
+
   /**
    * Maps over elements of the stream with the specified effectful function,
    * executing up to `n` invocations of `f` concurrently. Transformed elements
@@ -1798,6 +1817,49 @@ object ZPipeline extends ZPipelinePlatformSpecificConstructors {
     new ZPipeline(loop(Chunk.ChunkIterator.empty, 0))
   }
 
+  /**
+   * Creates a pipeline that maps over elements of the stream with the specified
+   * effectful function.
+   *
+   * Unlike `mapZIO` processing is done chunk by chunk. This means that
+   * `mapZIOChunked` provides weaker guarantees than `mapZIO`. While
+   * `stream.mapZIO(f).mapZIO(g)` is guaranteed to be equivalent to
+   * `stream.mapZIO(x => f(x).flatMap(g))`, the same is not true for
+   * `mapZIOChunked`. For example, `mapZIO` guarantees that the first element of
+   * a stream will first be processed with `f` and then `g` before the second
+   * element is processed with `f`. `mapZIOChunked` may process the first two
+   * elements with `f` and only then move on to process the first element with
+   * `g`.
+   */
+  def mapZIOChunked[Env, Err, In, Out](
+    f: In => ZIO[Env, Err, Out]
+  )(implicit trace: Trace): ZPipeline[Env, Err, In, Out] = {
+    def writeWithNext(
+      builder: ChunkBuilder[Out],
+      next: ZChannel[Env, Err, Chunk[In], Any, Err, Chunk[Out], Any]
+    ): ZChannel[Env, Err, Chunk[In], Any, Err, Chunk[Out], Any] = {
+      val out = builder.result()
+      if (out.nonEmpty) ZChannel.write(out) *> next else next
+    }
+
+    lazy val reader: ZChannel[Env, Err, Chunk[In], Any, Err, Chunk[Out], Any] =
+      ZChannel.readWithCause(
+        chunk =>
+          ZChannel.unwrap {
+            val builder = ChunkBuilder.make[Out](chunk.size)
+            chunk
+              .mapZIODiscard(f(_).map(builder += _))
+              .foldCause(
+                cause => writeWithNext(builder, ZChannel.refailCause(cause)),
+                _ => writeWithNext(builder, reader)
+              )
+          },
+        err => ZChannel.refailCause(err),
+        done => ZChannel.succeed(done)
+      )
+    new ZPipeline(reader)
+  }
+
   /**
    * Maps over elements of the stream with the specified effectful function,
    * executing up to `n` invocations of `f` concurrently. Transformed elements

diff --git a/streams/shared/src/main/scala/zio/stream/ZStream.scala b/streams/shared/src/main/scala/zio/stream/ZStream.scala
@@ -1922,6 +1922,8 @@ final class ZStream[-R, +E, +A] private (val channel: ZChannel[R, Any, Any, Any,
    *
    * @note
    *   This combinator destroys the chunking structure.
+   * @see
+   *   [[mapZIOChunked]] for a version that preserves the chunking structure
    */
   def mapZIO[R1 <: R, E1 >: E, A1](f: A => ZIO[R1, E1, A1])(implicit trace: Trace): ZStream[R1, E1, A1] = {
 
@@ -1943,6 +1945,23 @@ final class ZStream[-R, +E, +A] private (val channel: ZChannel[R, Any, Any, Any,
     new ZStream(self.channel >>> loop(Chunk.ChunkIterator.empty, 0))
   }
 
+  /**
+   * Creates a pipeline that maps over elements of the stream with the specified
+   * effectful function.
+   *
+   * Unlike `mapZIO` processing is done chunk by chunk. This means that
+   * `mapZIOChunked` provides weaker guarantees than `mapZIO`. While
+   * `stream.mapZIO(f).mapZIO(g)` is guaranteed to be equivalent to
+   * `stream.mapZIO(x => f(x).flatMap(g))`, the same is not true for
+   * `mapZIOChunked`. For example, `mapZIO` guarantees that the first element of
+   * a stream will first be processed with `f` and then `g` before the second
+   * element is processed with `f`. `mapZIOChunked` may process the first two
+   * elements with `f` and only then move on to process the first element with
+   * `g`.
+   */
+  def mapZIOChunked[R1 <: R, E1 >: E, A1](f: A => ZIO[R1, E1, A1])(implicit trace: Trace): ZStream[R1, E1, A1] =
+    self >>> ZPipeline.mapZIOChunked(f)
+
   /**
    * Maps over elements of the stream with the specified effectful function,
    * executing up to `n` invocations of `f` concurrently. Transformed elements