zio · adamgfraser · Jun 1, 2023 · Apr 18, 2023 · Apr 18, 2023 · Apr 19, 2023
diff --git a/streams-tests/shared/src/test/scala/zio/stream/ZPipelineSpec.scala b/streams-tests/shared/src/test/scala/zio/stream/ZPipelineSpec.scala
@@ -153,6 +153,63 @@ object ZPipelineSpec extends ZIOBaseSpec {
             .exit
         )(fails(equalTo("failed!!!")))
       ),
+      suite("sample")(
+        test("Works with empty input") {
+          for {
+            result <- ZStream.empty.via(ZPipeline.sample(0.5d)).run(ZSink.collectAll)
+          } yield assert(result)(equalTo(Chunk.empty[Any]))
+        },
+        test("Keeps everything when passed 1.0d") {
+          val range = 1.to(10000)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.sample(1.0d))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(equalTo(chunk))
+        },
+        test("Keeps nothing when passed 0.0d") {
+          val range = 1.to(10000)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.sample(0.0d))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(equalTo(Chunk.empty[Int]))
+        },
+        test("Keeps about half when passed 0.5d") {
+          val range = 1.to(10000)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.sample(0.5d))
+                        .run(ZSink.collectAll)
+          } yield assert(result.size)(equalTo(5062))
+        },
+        test("Keeps about a quarter when passed 0.25d") {
+          val range = 1.to(10000)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.sample(0.25d))
+                        .run(ZSink.collectAll)
+          } yield assert(result.size)(equalTo(2543))
+        },
+        test("Keeps order and values intact") {
+          val range = 1.to(100)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.sample(0.05d))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(equalTo(Chunk(3, 5, 10, 36, 54, 60, 80)))
+        }
+      ),
       suite("hex")(
         test("Empty input encodes to empty output") {
           for {

diff --git a/streams/shared/src/main/scala/zio/stream/ZPipeline.scala b/streams/shared/src/main/scala/zio/stream/ZPipeline.scala
@@ -1826,6 +1826,18 @@ object ZPipeline extends ZPipelinePlatformSpecificConstructors {
     new ZPipeline(ZChannel.suspend(process(new ZStream.Rechunker(target), target)))
   }
 
+  /**
+   * Creates a pipeline that randomly samples elements according to the
+   * specified percentage.
+   */
+  def sample[In](p: => Double)(implicit trace: Trace): ZPipeline[Any, Nothing, In, In] = {
+    val clamped = if (p.isNaN || p < 0.0d) 0.0d else if (p > 1.0d) 1.0d else p
+    val channel = ZChannel
+      .identity[Nothing, Chunk[In], Any]
+      .mapOutZIO(_.filterZIO(_ => Random.nextDoubleBetween(0.0d, 1.0d).map(_ < clamped)))
+    ZPipeline.fromChannel(channel)
+  }
+
   /**
    * Creates a pipeline that scans elements with the specified function.
    */