drcrallen commented on a change in pull request #7174: Change reservoir
sampling of segments to spliterator approach
URL: https://github.com/apache/incubator-druid/pull/7174#discussion_r262694458
##########
File path:
server/src/main/java/org/apache/druid/server/coordinator/ReservoirSegmentSampler.java
##########
@@ -19,36 +19,93 @@
package org.apache.druid.server.coordinator;
+import com.google.common.math.IntMath;
+import org.apache.druid.client.ImmutableDruidServer;
import org.apache.druid.timeline.DataSegment;
+import java.math.RoundingMode;
+import java.util.Collection;
import java.util.List;
+import java.util.Spliterator;
import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
final class ReservoirSegmentSampler
{
+ private static final int SPLITERATOR_SIZE_THRESHOLD = 25;
static BalancerSegmentHolder getRandomBalancerSegmentHolder(final
List<ServerHolder> serverHolders)
{
- ServerHolder fromServerHolder = null;
+ ServerHolder fromServerHolder = reservoirSampleServer(serverHolders);
+ if (fromServerHolder == null) {
+ return null;
+ }
+
DataSegment proposalSegment = null;
- int numSoFar = 0;
-
- for (ServerHolder server : serverHolders) {
- for (DataSegment segment : server.getServer().getSegments()) {
- int randNum = ThreadLocalRandom.current().nextInt(numSoFar + 1);
- // w.p. 1 / (numSoFar+1), swap out the server and segment
- if (randNum == numSoFar) {
- fromServerHolder = server;
- proposalSegment = segment;
+ Collection<DataSegment> segments =
fromServerHolder.getServer().getSegments();
+
+ if (segments.size() == 1) {
+ proposalSegment = segments.iterator().next();
+ } else {
+ Spliterator<DataSegment> chosenSpliterator = segments.spliterator();
+ if (chosenSpliterator == null) {
+ return null;
+ }
+
+ for (int i = 0; i < IntMath.log2(segments.size(), RoundingMode.UP); i++)
{
+ if (chosenSpliterator.estimateSize() < SPLITERATOR_SIZE_THRESHOLD) {
+ List<DataSegment> finalList = StreamSupport
+ .stream(chosenSpliterator, false)
+ .collect(Collectors.toList());
+
+ proposalSegment = getRandomElementFromList(finalList);
+ break;
+ }
+
+ Spliterator<DataSegment> newSpliterator = chosenSpliterator.trySplit();
+
+ // Choose between itself or the new spliterator with equal probability
+ if (ThreadLocalRandom.current().nextBoolean()) {
+ chosenSpliterator = newSpliterator;
}
- numSoFar++;
}
}
- if (fromServerHolder != null) {
- return new BalancerSegmentHolder(fromServerHolder.getServer(),
proposalSegment);
- } else {
+ return new BalancerSegmentHolder(fromServerHolder.getServer(),
proposalSegment);
+ }
+
+ static ServerHolder reservoirSampleServer(final List<ServerHolder>
serverHolders)
+ {
+ if (serverHolders.isEmpty()) {
+ return null;
+ }
+
+ ServerHolder sampledServer = null;
+ long total = 0;
+
+ for (ServerHolder serverHolder : serverHolders) {
+ ImmutableDruidServer server = serverHolder.getServer();
+ long numSegments = server.getSegments().size();
+ long endIndex = total + numSegments;
+ // Handle edge case where first server contains no segments so that
nextLong() doesn't fail
+ long randomIndex = ThreadLocalRandom.current().nextLong((endIndex == 0)
? 1 : endIndex);
Review comment:
can you split this out explicitly in an if statement rather than having
`nextLong(1)` as a nested consideration
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
