[GitHub] [druid] jihoonson commented on a change in pull request #11257: Reduce method invocation of reservoir sampling

[GitBox <gi...@apache.org>]

jihoonson commented on a change in pull request #11257:
URL: https://github.com/apache/druid/pull/11257#discussion_r678820747



##########
File path: server/src/main/java/org/apache/druid/server/coordinator/BalancerStrategy.java
##########
@@ -55,28 +55,90 @@
   ServerHolder findNewSegmentHomeReplicator(DataSegment proposalSegment, List<ServerHolder> serverHolders);
 
   /**
-   * Pick the best segment to move from one of the supplied set of servers according to the balancing strategy.
+   * Pick the best segments to move from one of the supplied set of servers according to the balancing strategy.
+   *
    * @param serverHolders set of historicals to consider for moving segments
    * @param broadcastDatasources Datasources that contain segments which were loaded via broadcast rules.
    *                             Balancing strategies should avoid rebalancing segments for such datasources, since
    *                             they should be loaded on all servers anyway.
    *                             NOTE: this should really be handled on a per-segment basis, to properly support
    *                                   the interval or period-based broadcast rules. For simplicity of the initial
    *                                   implementation, only forever broadcast rules are supported.
+   * @param reservoirSize the reservoir size maintained by the Reservoir Sampling algorithm.
    * @param percentOfSegmentsToConsider The percentage of the total number of segments that we will consider when
    *                                    choosing which segment to move. {@link CoordinatorDynamicConfig} defines a
    *                                    config percentOfSegmentsToConsiderPerMove that will be used as an argument
    *                                    for implementations of this method.
-   *
-   * @return {@link BalancerSegmentHolder} containing segment to move and server it currently resides on, or null if
-   *         there are no segments to pick from (i. e. all provided serverHolders are empty).
+   * @return Iterator for set of {@link BalancerSegmentHolder} containing segment to move and server they currently
+   * reside on, or empty if there are no segments to pick from (i. e. all provided serverHolders are empty).
    */
-  @Nullable
-  BalancerSegmentHolder pickSegmentToMove(
+  default Iterator<BalancerSegmentHolder> pickSegmentsToMove(
       List<ServerHolder> serverHolders,
       Set<String> broadcastDatasources,
+      int reservoirSize,
       double percentOfSegmentsToConsider
-  );
+  )
+  {
+    if (reservoirSize > 1) {
+      return new Iterator<BalancerSegmentHolder>()
+      {
+        private Iterator<BalancerSegmentHolder> it = sample();
+
+        private Iterator<BalancerSegmentHolder> sample()
+        {
+          return ReservoirSegmentSampler.getRandomBalancerSegmentHolders(
+              serverHolders,
+              broadcastDatasources,
+              reservoirSize
+          ).iterator();
+        }
+
+        @Override
+        public boolean hasNext()
+        {
+          if (it.hasNext()) {
+            return true;
+          }
+          it = sample();
+          return it.hasNext();
+        }
+
+        @Override
+        public BalancerSegmentHolder next()
+        {
+          return it.next();
+        }
+      };
+    }
+
+    return new Iterator<BalancerSegmentHolder>()
+    {
+      private BalancerSegmentHolder next = sample();
+
+      private BalancerSegmentHolder sample()
+      {
+        return ReservoirSegmentSampler.getRandomBalancerSegmentHolder(
+            serverHolders,
+            broadcastDatasources,
+            percentOfSegmentsToConsider
+        );
+      }
+
+      @Override
+      public boolean hasNext()
+      {
+        return next != null;
+      }
+
+      @Override
+      public BalancerSegmentHolder next()
+      {

Review comment:
       nit:
   
   ```java
           if (!hasNext()) {
             throw new NoSuchElementException();
           }
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org