Optimize ZQuery.collectAllPar For Very Small Collections (#349)

* optimize * format
zio · Jul 28, 2022 · d707caa · d707caa
1 parent d2812dd
commit d707caa
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/zio-query/shared/src/main/scala/zio/query/ZQuery.scala b/zio-query/shared/src/main/scala/zio/query/ZQuery.scala
@@ -1118,7 +1118,15 @@ object ZQuery {
   )(
     f: A => ZQuery[R, E, B]
   )(implicit bf: BuildFrom[Collection[A], B, Collection[B]], trace: Trace): ZQuery[R, E, Collection[B]] =
-    ZQuery(ZIO.foreachPar(Chunk.fromIterable(as))(f(_).step).map(Result.collectAllPar(_).map(bf.fromSpecific(as))))
+    ZQuery.suspend {
+      val size = as.size
+      if (size == 0)
+        ZQuery.succeed(bf.newBuilder(as).result())
+      else if (size == 1)
+        f(as.head).map(bf.newBuilder(as) += _).map(_.result())
+      else
+        ZQuery(ZIO.foreachPar(Chunk.fromIterable(as))(f(_).step).map(Result.collectAllPar(_).map(bf.fromSpecific(as))))
+    }
 
   /**
    * Performs a query for each element in a Set, collecting the results