Merge pull request #93 from betarixm/fix/setup-partition-max-size

Set max size of sened partition
betarixm · Nov 29, 2023 · 41006c4 · 41006c4
2 parents 34d543c + 82c30fa
commit 41006c4
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 15 deletions.
diff --git a/core/src/main/scala/Block.scala b/core/src/main/scala/Block.scala
@@ -91,4 +91,8 @@ class Block(val records: LazyList[Record]) extends AnyVal {
   def sample(number: Int = 64): LazyList[Key] =
     Record.sample(records, number)
 
+  def splitIntoChunks(chunkSize: Int): LazyList[Block] = {
+    records.grouped(chunkSize).map(new Block(_)).to(LazyList)
+  }
+
 }
diff --git a/master/src/main/scala/Master.scala b/master/src/main/scala/Master.scala
@@ -66,7 +66,7 @@ object Master extends Logging {
     logger.info("[Master] Sample Requested")
 
     val sampledKeys = clients
-      .sample(64)(requestExecutionContext)
+      .sample(1024)(requestExecutionContext)
       .flatMap(_.sampledKeys)
       .map(Key.fromByteString)
 

diff --git a/worker/src/main/scala/WorkerService.scala b/worker/src/main/scala/WorkerService.scala
@@ -200,21 +200,25 @@ class WorkerService(
         inputFiles.foreach(path => {
           val block = Block.fromPath(path)
 
-          val targetClients = workersWithClients
-            .filter(
-              _._1.keyRange.get.includes(block.records.head.key)
+          val blockChunks = block.splitIntoChunks(1000)
+
+          blockChunks.foreach(chunk => {
+            val targetClients = workersWithClients
+              .filter(
+                _._1.keyRange.get.includes(chunk.records.head.key)
+              )
+              .map(_._2)
+
+            logger.info(s"[WorkerServer] Sending $chunk to $targetClients")
+
+            Await.result(
+              Future.traverse(targetClients)(_.saveBlock(chunk))(
+                GenericBuildFrom[WorkerClient, SaveBlockReply],
+                executionContext
+              ),
+              scala.concurrent.duration.Duration.Inf
             )
-            .map(_._2)
-
-          logger.info(s"[WorkerServer] Sending $block to $targetClients")
-
-          Await.result(
-            Future.traverse(targetClients)(_.saveBlock(block))(
-              GenericBuildFrom[WorkerClient, SaveBlockReply],
-              executionContext
-            ),
-            scala.concurrent.duration.Duration.Inf
-          )
+          })
         })
 
         clients.foreach(_.shutdown())