improve performance of parallel operations (#90)

zio · Sep 6, 2020 · 77a694d · 77a694d
1 parent 29c6192
commit 77a694d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 11 deletions.
diff --git a/zio-query/shared/src/main/scala/zio/query/ZQuery.scala b/zio-query/shared/src/main/scala/zio/query/ZQuery.scala
@@ -442,10 +442,17 @@ object ZQuery {
   def foreach[R, E, A, B, Collection[+Element] <: Iterable[Element]](
     as: Collection[A]
   )(f: A => ZQuery[R, E, B])(implicit bf: BuildFrom[Collection[A], B, Collection[B]]): ZQuery[R, E, Collection[B]] =
-    as.foldLeft[ZQuery[R, E, Builder[B, Collection[B]]]](ZQuery.succeed(bf.newBuilder(as)))((bs, a) =>
-        bs.zipWith(f(a))(_ += _)
-      )
-      .map(_.result())
+    if (as.isEmpty) ZQuery.succeed(bf.newBuilder(as).result)
+    else {
+      val iterator                                         = as.iterator
+      var builder: ZQuery[R, E, Builder[B, Collection[B]]] = null
+      while (iterator.hasNext) {
+        val a = iterator.next()
+        if (builder eq null) builder = f(a).map(bf.newBuilder(as) += _)
+        else builder = builder.zipWith(f(a))(_ += _)
+      }
+      builder.map(_.result())
+    }
 
   /**
    * Performs a query for each element in a collection, collecting the results
@@ -455,10 +462,17 @@ object ZQuery {
   def foreachPar[R, E, A, B, Collection[+Element] <: Iterable[Element]](
     as: Collection[A]
   )(f: A => ZQuery[R, E, B])(implicit bf: BuildFrom[Collection[A], B, Collection[B]]): ZQuery[R, E, Collection[B]] =
-    as.foldLeft[ZQuery[R, E, Builder[B, Collection[B]]]](ZQuery.succeed(bf.newBuilder(as)))((bs, a) =>
-        bs.zipWithPar(f(a))(_ += _)
-      )
-      .map(_.result())
+    if (as.isEmpty) ZQuery.succeed(bf.newBuilder(as).result)
+    else {
+      val iterator                                         = as.iterator
+      var builder: ZQuery[R, E, Builder[B, Collection[B]]] = null
+      while (iterator.hasNext) {
+        val a = iterator.next()
+        if (builder eq null) builder = f(a).map(bf.newBuilder(as) += _)
+        else builder = builder.zipWithPar(f(a))(_ += _)
+      }
+      builder.map(_.result())
+    }
 
   /**
    * Constructs a query from an effect.

diff --git a/zio-query/shared/src/main/scala/zio/query/internal/Continue.scala b/zio-query/shared/src/main/scala/zio/query/internal/Continue.scala
@@ -114,9 +114,9 @@ private[query] sealed trait Continue[-R, +E, +A] { self =>
   final def zipWithPar[R1 <: R, E1 >: E, B, C](that: Continue[R1, E1, B])(f: (A, B) => C): Continue[R1, E1, C] =
     (self, that) match {
       case (Effect(l), Effect(r)) => effect(l.zipWithPar(r)(f))
-      case (Effect(l), Get(r))    => effect(l.zipWithPar(ZQuery.fromEffect(r))(f))
-      case (Get(l), Effect(r))    => effect(ZQuery.fromEffect(l).zipWithPar(r)(f))
-      case (Get(l), Get(r))       => get(l.zipWithPar(r)(f))
+      case (Effect(l), Get(r))    => effect(l.zipWith(ZQuery.fromEffect(r))(f))
+      case (Get(l), Effect(r))    => effect(ZQuery.fromEffect(l).zipWith(r)(f))
+      case (Get(l), Get(r))       => get(l.zipWith(r)(f))
     }
 
 }