clear precomputedSplitAssignments

Signed-off-by: Tilak Raj <[email protected]>
synadia-io · Jan 24, 2025 · 054127c · 054127c
1 parent 6ba9b30
commit 054127c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 16 deletions.
diff --git a/src/main/java/io/synadia/flink/v0/enumerator/NatsSourceEnumerator.java b/src/main/java/io/synadia/flink/v0/enumerator/NatsSourceEnumerator.java
@@ -63,22 +63,17 @@ private List<List<NatsSubjectSplit>> preComputeSplitsAssignments (int parallelis
 
         // Distribute splits evenly among subtasks
         for (int j = 0; j < parallelism; j++) {
-            for (int i = 0; i < minimumSplitsPerReader; i++) {
-                NatsSubjectSplit split = remainingSplits.poll();
-                if (split != null) {
-                    splitAssignments.get(j).add(split);
-                }
+            List<NatsSubjectSplit> readerSplits = splitAssignments.get(j);
+
+            // Assign minimum splits to each reader
+            for (int i = 0; i < minimumSplitsPerReader && !remainingSplits.isEmpty(); i++) {
+                readerSplits.add(remainingSplits.poll());
             }
 
             // Assign one leftover split if available
             if (leftoverSplits > 0 && !remainingSplits.isEmpty()) {
-                NatsSubjectSplit extraSplit = remainingSplits.poll();
-                if (extraSplit != null) {
-                    splitAssignments.get(j).add(extraSplit);
-
-                    // Reduce leftover count
-                    leftoverSplits--;
-                }
+                readerSplits.add(remainingSplits.poll());
+                leftoverSplits--;
             }
         }
 
@@ -87,6 +82,8 @@ private List<List<NatsSubjectSplit>> preComputeSplitsAssignments (int parallelis
 
     @Override
     public void close() {
+        // remove precomputed split assignments if any
+        precomputedSplitAssignments.clear();
     }
 
     @Override
@@ -98,6 +95,8 @@ public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname
             context.signalNoMoreSplits(subtaskId);
         } else {
             // O(1) operation with LinkedList
+            // Remove the first element from the list
+            // and assign splits to subtask
             List<NatsSubjectSplit> splits = precomputedSplitAssignments.remove(0);
             if (splits.isEmpty()) {
                 LOG.debug("{} | Empty split assignment for subtask {}", id, subtaskId);

diff --git a/src/test/java/io/synadia/io/synadia/flink/v0/NatsSourceEnumeratorTests.java b/src/test/java/io/synadia/io/synadia/flink/v0/NatsSourceEnumeratorTests.java
@@ -103,7 +103,6 @@ void testHandleSplitRequestWithSplitsGreaterThanParallelism() {
             });
         }
 
-        // Expected: Each subtask gets at least 2 splits (6 splits, 3 subtasks)
         assertEquals(5, assignedSplits.size());
         assertEquals(List.of("split1", "split2", "split3", "split4", "split5"), assignedSplits);
     }
@@ -143,7 +142,6 @@ void testHandleSplitRequestWithSplitsEqualToParallelism() {
             });
         }
 
-        // Expected: Each subtask gets at least 2 splits (6 splits, 3 subtasks)
         assertEquals(3, assignedSplits.size());
         assertEquals(List.of("split1", "split2", "split3"), assignedSplits);
     }
@@ -188,7 +186,6 @@ void testHandleSplitRequestWithLargeParallelism() {
             });
         }
 
-        // Expected: Each subtask gets at least 2 splits (6 splits, 3 subtasks)
         assertEquals(9, assignedSplits.size());
         assertEquals(List.of("split1", "split2", "split3", "split4", "split5", "split6", "split7", "split8", "split9"), assignedSplits);
     }
@@ -227,7 +224,6 @@ void testHandleSplitRequestWithLargeSplits() {
             });
         }
 
-        // Expected: Each subtask gets at least 2 splits (6 splits, 3 subtasks)
         assertEquals(100, assignedSplits.size());
     }
 }