Avoid filtering requests when CompletedRequestsMap is empty (#466)

* Optimize `DataSource.Batched.runAll` * Check if requests are not empty prior to running them
zio · Mar 29, 2024 · 170b25c · 170b25c
1 parent 99f5ac9
commit 170b25c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/zio-query/shared/src/main/scala/zio/query/CompletedRequestMap.scala b/zio-query/shared/src/main/scala/zio/query/CompletedRequestMap.scala
@@ -67,6 +67,12 @@ final class CompletedRequestMap private (private val map: Map[Any, Exit[Any, Any
   def requests: Set[Request[_, _]] =
     map.keySet.asInstanceOf[Set[Request[_, _]]]
 
+  /**
+   * Whether the completed requests map is empty.
+   */
+  def isEmpty: Boolean =
+    map.isEmpty
+
   override def toString: String =
     s"CompletedRequestMap(${map.mkString(", ")})"
 }

diff --git a/zio-query/shared/src/main/scala/zio/query/DataSource.scala b/zio-query/shared/src/main/scala/zio/query/DataSource.scala
@@ -171,9 +171,12 @@ object DataSource {
     def run(requests: Chunk[A])(implicit trace: Trace): ZIO[R, Nothing, CompletedRequestMap]
     final def runAll(requests: Chunk[Chunk[A]])(implicit trace: Trace): ZIO[R, Nothing, CompletedRequestMap] =
       ZIO.foldLeft(requests)(CompletedRequestMap.empty) { case (completedRequestMap, requests) =>
-        val newRequests = requests.filterNot(completedRequestMap.contains)
-        if (newRequests.isEmpty) ZIO.succeed(completedRequestMap)
-        else run(newRequests).map(completedRequestMap ++ _)
+        if (completedRequestMap.isEmpty && requests.nonEmpty) run(requests)
+        else {
+          val newRequests = requests.filterNot(completedRequestMap.contains)
+          if (newRequests.isEmpty) ZIO.succeed(completedRequestMap)
+          else run(newRequests).map(completedRequestMap ++ _)
+        }
       }
   }