Jackie-Jiang commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4342835555
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1622,174 +1656,280 @@ private static void handleErrorInstance(String
segmentName, String instanceName,
}
}
+ /// Returns the target assignment for the ideal state changed while waiting
for the external view to converge, when
+ /// the segments to be moved are unchanged: the current assignment with the
target of the segments to be moved. For
+ /// a partition-pinned table (strict replica-group routing), a partition
added while waiting keeps its current
+ /// instances, and a segment added while waiting, which is assigned to the
current instances of its partition, is
+ /// moved to the target instances of its partition, unless it is a CONSUMING
segment and CONSUMING segments are not
+ /// included in the rebalance. Returns `null` when the target assignment
must be recalculated instead: when the
+ /// partition id of such a segment is not in its name, or when CONSUMING
segments are not included and an ONLINE
+ /// segment was added to a pauseless table, because it can still be
committing (which only its ZK metadata tells),
+ /// in which case it is not moved either.
+ @Nullable
+ @VisibleForTesting
+ static TableAssignment getTargetAssignmentForNewIdealState(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, List<String>
segmentsToMove,
+ boolean enableStrictReplicaGroup, boolean includeConsuming, boolean
pauseless, Logger tableRebalanceLogger) {
+ // Other instance assignment code returns a TreeMap to keep it sorted,
doing the same here
+ Map<String, Map<String, String>> newTargetAssignment = new
TreeMap<>(currentAssignment);
+ for (String segment : segmentsToMove) {
+ newTargetAssignment.put(segment, targetAssignment.get(segment));
+ }
+ if (!enableStrictReplicaGroup) {
+ return new TableAssignment(newTargetAssignment,
targetPartitionAssignment);
+ }
+ Map<Integer, List<String>> newTargetPartitionAssignment = new
TreeMap<>(targetPartitionAssignment);
+ for (Map.Entry<Integer, List<String>> entry :
currentPartitionAssignment.entrySet()) {
+ newTargetPartitionAssignment.putIfAbsent(entry.getKey(),
entry.getValue());
Review Comment:
Fixed in 7df9c15: only the partitions added since the previous IdealState
read are re-added; partitions dropped by the target stay dropped. Covered in
`TableRebalancerTest`.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1622,174 +1656,280 @@ private static void handleErrorInstance(String
segmentName, String instanceName,
}
}
+ /// Returns the target assignment for the ideal state changed while waiting
for the external view to converge, when
+ /// the segments to be moved are unchanged: the current assignment with the
target of the segments to be moved. For
+ /// a partition-pinned table (strict replica-group routing), a partition
added while waiting keeps its current
+ /// instances, and a segment added while waiting, which is assigned to the
current instances of its partition, is
+ /// moved to the target instances of its partition, unless it is a CONSUMING
segment and CONSUMING segments are not
+ /// included in the rebalance. Returns `null` when the target assignment
must be recalculated instead: when the
+ /// partition id of such a segment is not in its name, or when CONSUMING
segments are not included and an ONLINE
+ /// segment was added to a pauseless table, because it can still be
committing (which only its ZK metadata tells),
+ /// in which case it is not moved either.
+ @Nullable
+ @VisibleForTesting
+ static TableAssignment getTargetAssignmentForNewIdealState(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, List<String>
segmentsToMove,
+ boolean enableStrictReplicaGroup, boolean includeConsuming, boolean
pauseless, Logger tableRebalanceLogger) {
+ // Other instance assignment code returns a TreeMap to keep it sorted,
doing the same here
+ Map<String, Map<String, String>> newTargetAssignment = new
TreeMap<>(currentAssignment);
+ for (String segment : segmentsToMove) {
+ newTargetAssignment.put(segment, targetAssignment.get(segment));
+ }
+ if (!enableStrictReplicaGroup) {
+ return new TableAssignment(newTargetAssignment,
targetPartitionAssignment);
+ }
+ Map<Integer, List<String>> newTargetPartitionAssignment = new
TreeMap<>(targetPartitionAssignment);
+ for (Map.Entry<Integer, List<String>> entry :
currentPartitionAssignment.entrySet()) {
+ newTargetPartitionAssignment.putIfAbsent(entry.getKey(),
entry.getValue());
+ }
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segment = entry.getKey();
+ if (targetAssignment.containsKey(segment)) {
+ continue;
+ }
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ String state;
+ if (currentInstanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
+ if (!includeConsuming && pauseless) {
+ tableRebalanceLogger.info("Found new ONLINE segment: {} of pauseless
table while CONSUMING segments are not "
+ + "included, re-calculating the target assignment based on the
new IdealState", segment);
+ return null;
+ }
+ state = SegmentStateModel.ONLINE;
+ } else if
(currentInstanceStateMap.containsValue(SegmentStateModel.CONSUMING)) {
+ if (!includeConsuming) {
+ continue;
+ }
+ state = SegmentStateModel.CONSUMING;
+ } else {
+ // All instances are OFFLINE, and OFFLINE segments are not rebalanced
+ continue;
+ }
+ Integer partitionId =
SegmentUtils.getPartitionIdFromSegmentName(segment);
+ if (partitionId == null) {
+ tableRebalanceLogger.info("Found new segment: {} without partition id
in its name, re-calculating the target "
+ + "assignment based on the new IdealState", segment);
+ return null;
+ }
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ if (currentInstances == null || currentInstances.size() !=
currentInstanceStateMap.size()
+ || !currentInstanceStateMap.keySet().containsAll(currentInstances)) {
+ // Not assigned to the instances of its partition (e.g. relocated to
the COMPLETED instances), keep it as is
+ continue;
+ }
+ newTargetAssignment.put(segment,
Review Comment:
Fixed in 7df9c15: a new completed segment of a table relocating completed
segments triggers a full recalculation instead of being re-targeted to the
CONSUMING placement.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1622,174 +1656,280 @@ private static void handleErrorInstance(String
segmentName, String instanceName,
}
}
+ /// Returns the target assignment for the ideal state changed while waiting
for the external view to converge, when
+ /// the segments to be moved are unchanged: the current assignment with the
target of the segments to be moved. For
+ /// a partition-pinned table (strict replica-group routing), a partition
added while waiting keeps its current
+ /// instances, and a segment added while waiting, which is assigned to the
current instances of its partition, is
+ /// moved to the target instances of its partition, unless it is a CONSUMING
segment and CONSUMING segments are not
+ /// included in the rebalance. Returns `null` when the target assignment
must be recalculated instead: when the
+ /// partition id of such a segment is not in its name, or when CONSUMING
segments are not included and an ONLINE
+ /// segment was added to a pauseless table, because it can still be
committing (which only its ZK metadata tells),
+ /// in which case it is not moved either.
+ @Nullable
+ @VisibleForTesting
+ static TableAssignment getTargetAssignmentForNewIdealState(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, List<String>
segmentsToMove,
+ boolean enableStrictReplicaGroup, boolean includeConsuming, boolean
pauseless, Logger tableRebalanceLogger) {
+ // Other instance assignment code returns a TreeMap to keep it sorted,
doing the same here
+ Map<String, Map<String, String>> newTargetAssignment = new
TreeMap<>(currentAssignment);
+ for (String segment : segmentsToMove) {
+ newTargetAssignment.put(segment, targetAssignment.get(segment));
+ }
+ if (!enableStrictReplicaGroup) {
+ return new TableAssignment(newTargetAssignment,
targetPartitionAssignment);
+ }
+ Map<Integer, List<String>> newTargetPartitionAssignment = new
TreeMap<>(targetPartitionAssignment);
+ for (Map.Entry<Integer, List<String>> entry :
currentPartitionAssignment.entrySet()) {
+ newTargetPartitionAssignment.putIfAbsent(entry.getKey(),
entry.getValue());
+ }
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segment = entry.getKey();
+ if (targetAssignment.containsKey(segment)) {
+ continue;
+ }
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ String state;
+ if (currentInstanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
+ if (!includeConsuming && pauseless) {
+ tableRebalanceLogger.info("Found new ONLINE segment: {} of pauseless
table while CONSUMING segments are not "
+ + "included, re-calculating the target assignment based on the
new IdealState", segment);
+ return null;
+ }
+ state = SegmentStateModel.ONLINE;
+ } else if
(currentInstanceStateMap.containsValue(SegmentStateModel.CONSUMING)) {
+ if (!includeConsuming) {
+ continue;
+ }
+ state = SegmentStateModel.CONSUMING;
+ } else {
+ // All instances are OFFLINE, and OFFLINE segments are not rebalanced
+ continue;
+ }
+ Integer partitionId =
SegmentUtils.getPartitionIdFromSegmentName(segment);
+ if (partitionId == null) {
+ tableRebalanceLogger.info("Found new segment: {} without partition id
in its name, re-calculating the target "
+ + "assignment based on the new IdealState", segment);
+ return null;
+ }
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ if (currentInstances == null || currentInstances.size() !=
currentInstanceStateMap.size()
+ || !currentInstanceStateMap.keySet().containsAll(currentInstances)) {
+ // Not assigned to the instances of its partition (e.g. relocated to
the COMPLETED instances), keep it as is
+ continue;
+ }
+ newTargetAssignment.put(segment,
+
SegmentAssignmentUtils.getInstanceStateMap(newTargetPartitionAssignment.get(partitionId),
state));
+ }
+ return new TableAssignment(newTargetAssignment,
newTargetPartitionAssignment);
+ }
+
/// Uses the default LOGGER
@VisibleForTesting
static Map<String, Map<String, String>> getNextAssignment(Map<String,
Map<String, String>> currentAssignment,
Map<String, Map<String, String>> targetAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
- boolean lowDiskMode, int batchSizePerServer,
Object2IntOpenHashMap<String> segmentPartitionIdMap,
- PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor
dataLossRiskAssessor) {
- return getNextAssignment(currentAssignment, targetAssignment,
minAvailableReplicas, enableStrictReplicaGroup,
- lowDiskMode, batchSizePerServer, segmentPartitionIdMap,
partitionIdFetcher, dataLossRiskAssessor, LOGGER);
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor) {
+ return getNextAssignment(currentAssignment, targetAssignment, Map.of(),
Map.of(), minAvailableReplicas,
+ enableStrictReplicaGroup, lowDiskMode, batchSizePerServer,
dataLossRiskAssessor).getSegmentAssignment();
+ }
+
+ /// Uses the default LOGGER
+ @VisibleForTesting
+ static TableAssignment getNextAssignment(Map<String, Map<String, String>>
currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor) {
+ return getNextAssignment(currentAssignment, targetAssignment,
currentPartitionAssignment,
+ targetPartitionAssignment, minAvailableReplicas,
enableStrictReplicaGroup, lowDiskMode, batchSizePerServer,
+ dataLossRiskAssessor, LOGGER);
}
/// Returns the next assignment for the table based on the current
assignment and the target assignment with regard to
/// the minimum available replicas requirement. For strict replica-group
mode, track the available instances for all
/// the segments with the same instances in the next assignment, and ensure
the minimum available replicas requirement
/// is met. If adding the assignment for a segment breaks the requirement,
use the current assignment for the segment.
///
- /// For strict replica group routing only (where the segment assignment is
not StrictRealtimeSegmentAssignment)
- /// if batching is enabled, the instances assigned for the same partitionId
can be different for different segments.
- /// For strict replica group routing with StrictRealtimeSegmentAssignment on
the other hand, the assignment for a
- /// given
- /// partitionId will be the same across all segments. We can treat both
cases similarly by creating a mapping from
- /// partitionId -> unique set of instance assignments -> currentAssignment.
With StrictRealtimeSegmentAssignment,
- /// this map will have a single entry for 'unique set of instance
assignments'.
- ///
- /// TODO: Ideally if strict replica group routing is enabled then
StrictRealtimeSegmentAssignment should be used, but
- /// this is not enforced in the code today. Once enforcement is added,
then the routing side and assignment side
- /// will be equivalent and all segments belonging to a given
partitionId will be assigned to the same set of
- /// instances. Special handling to check each group of assigned
instances can be removed in that case. The
- /// strict replica group routing can also be utilized for OFFLINE
tables, thus StrictRealtimeSegmentAssignment
- /// also needs to be made more generic for the OFFLINE case.
- private static Map<String, Map<String, String>>
getNextAssignment(Map<String, Map<String, String>> currentAssignment,
- Map<String, Map<String, String>> targetAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
- boolean lowDiskMode, int batchSizePerServer,
Object2IntOpenHashMap<String> segmentPartitionIdMap,
- PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor
dataLossRiskAssessor, Logger tableRebalanceLogger) {
+ /// The partition assignment (the instances of each partition of a
partition-pinned table, see
+ /// [PartitionAssignmentUtils]) is moved along with the segments: a
partition takes the same step as a segment with
+ /// the same current and target instances, so that a segment added to the
partition during the rebalance is assigned
+ /// to the same instances as the other segments of the partition at that
point. In strict replica-group mode, a
+ /// partition moves exactly when the segments with its current and target
instances move. A partition of the target
+ /// without entry in the current partition assignment is left without entry,
as its current instances are unknown:
+ /// the caller sets it once the segments reach the target. A partition of
the current partition assignment absent
+ /// from the target is dropped.
+ private static TableAssignment getNextAssignment(Map<String, Map<String,
String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor,
+ Logger tableRebalanceLogger) {
return enableStrictReplicaGroup ?
getNextStrictReplicaGroupAssignment(currentAssignment, targetAssignment,
- minAvailableReplicas, lowDiskMode, batchSizePerServer,
segmentPartitionIdMap, partitionIdFetcher,
+ currentPartitionAssignment, targetPartitionAssignment,
minAvailableReplicas, lowDiskMode, batchSizePerServer,
dataLossRiskAssessor, tableRebalanceLogger)
- : getNextNonStrictReplicaGroupAssignment(currentAssignment,
targetAssignment, minAvailableReplicas, lowDiskMode,
- batchSizePerServer, dataLossRiskAssessor);
+ : getNextNonStrictReplicaGroupAssignment(currentAssignment,
targetAssignment, currentPartitionAssignment,
+ targetPartitionAssignment, minAvailableReplicas, lowDiskMode,
batchSizePerServer, dataLossRiskAssessor);
}
- private static Map<String, Map<String, String>>
getNextStrictReplicaGroupAssignment(
+ private static TableAssignment getNextStrictReplicaGroupAssignment(
Map<String, Map<String, String>> currentAssignment, Map<String,
Map<String, String>> targetAssignment,
- int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer,
- Object2IntOpenHashMap<String> segmentPartitionIdMap, PartitionIdFetcher
partitionIdFetcher,
- DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) {
+ Map<Integer, List<String>> currentPartitionAssignment, Map<Integer,
List<String>> targetPartitionAssignment,
+ int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer,
DataLossRiskAssessor dataLossRiskAssessor,
+ Logger tableRebalanceLogger) {
Map<String, Map<String, String>> nextAssignment = new TreeMap<>();
+ Map<Integer, List<String>> nextPartitionAssignment = new TreeMap<>();
Map<String, Integer> numSegmentsToOffloadMap =
getNumSegmentsToOffloadMap(currentAssignment, targetAssignment);
Map<Pair<Set<String>, Set<String>>, Set<String>> assignmentMap = new
HashMap<>();
Map<Set<String>, Set<String>> availableInstancesMap = new HashMap<>();
Map<String, Integer> serverToNumSegmentsAddedSoFar = new HashMap<>();
if (batchSizePerServer == RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER) {
// Directly update the nextAssignment with anyServerExhaustedBatchSize =
false and return if batching is disabled
- updateNextAssignmentForPartitionIdStrictReplicaGroup(currentAssignment,
targetAssignment, nextAssignment, false,
- minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap, availableInstancesMap,
- serverToNumSegmentsAddedSoFar, dataLossRiskAssessor);
- return nextAssignment;
- }
-
- // Batching is enabled, calculate the Pair(current instances, target
instances) -> partitionId -> currentAssignment
- Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String,
String>>>>
- currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap =
-
getCurrentAndTargetInstancesToPartitionIdToCurrentAssignmentMap(currentAssignment,
targetAssignment,
- segmentPartitionIdMap, partitionIdFetcher);
-
- // Iterating over the unique pairs of current and target instances
- for (Map<Integer, Map<String, Map<String, String>>>
partitionIdToCurrentAssignment
- :
currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap.values()) {
- // Check if the servers of the first assignment for each unique
partition has any space left to move the
- // segments assigned to the partition and unique assigned instances as a
whole. If so, let's mark the partitions
- // as to be moved, otherwise we mark the partition as a whole as not
moveable.
- // Iterating over the partitionIds with the same unique pair of current
and assigned instances
- for (Map<String, Map<String, String>> curAssignment :
partitionIdToCurrentAssignment.values()) {
- Map.Entry<String, Map<String, String>> firstEntry =
curAssignment.entrySet().iterator().next();
- // It is enough to check for whether any server for one segment is
above the limit or not since all segments
- // in curAssignment will have the same current and target instances
and same partitionId
- Map<String, String> firstEntryInstanceStateMap = firstEntry.getValue();
- SingleSegmentAssignment firstAssignment =
- getNextSingleSegmentAssignment(firstEntryInstanceStateMap,
targetAssignment.get(firstEntry.getKey()),
- minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap);
- Set<String> serversAdded =
-
getServersAddedInSingleSegmentAssignment(firstEntryInstanceStateMap,
firstAssignment._instanceStateMap);
- boolean anyServerExhaustedBatchSize = false;
- for (String server : serversAdded) {
- int segmentsAddedToServerSoFar =
serverToNumSegmentsAddedSoFar.getOrDefault(server, 0);
- // Case I: We already exceeded the batchSizePerServer for this
server, cannot add any more segments
- // Case II: We have not yet exceeded the batchSizePerServer for this
server, but we don't have sufficient
- // space to host the segments for this assignment on the server, and
we have allocated some partitions so
- // far. If the batchSizePerServer is less than the number of
segments in a given partitionId, we must host
- // at least 1 partition and exceed the batchSizePerServer to ensure
progress is made. Thus, performing this
- // check only if segmentsAddedToServerSoFar > 0 is necessary.
- if ((segmentsAddedToServerSoFar >= batchSizePerServer) ||
(segmentsAddedToServerSoFar > 0
- && (segmentsAddedToServerSoFar + curAssignment.size()) >
batchSizePerServer)) {
- anyServerExhaustedBatchSize = true;
- break;
- }
+ StrictReplicaGroup group = new StrictReplicaGroup();
+ group._segments.putAll(currentAssignment);
+ for (int partitionId : currentPartitionAssignment.keySet()) {
+ if (targetPartitionAssignment.containsKey(partitionId)) {
+ group._partitionIds.add(partitionId);
}
- updateNextAssignmentForPartitionIdStrictReplicaGroup(curAssignment,
targetAssignment, nextAssignment,
- anyServerExhaustedBatchSize, minAvailableReplicas, lowDiskMode,
numSegmentsToOffloadMap, assignmentMap,
- availableInstancesMap, serverToNumSegmentsAddedSoFar,
dataLossRiskAssessor);
}
+ moveStrictReplicaGroup(group, false, targetAssignment,
currentPartitionAssignment, targetPartitionAssignment,
+ nextAssignment, nextPartitionAssignment, minAvailableReplicas,
lowDiskMode, numSegmentsToOffloadMap,
+ assignmentMap, availableInstancesMap, serverToNumSegmentsAddedSoFar,
dataLossRiskAssessor);
+ return new TableAssignment(nextAssignment, nextPartitionAssignment);
+ }
+
+ // Batching is enabled, move the segments and the partitions group by group
+ for (StrictReplicaGroup group : getStrictReplicaGroups(currentAssignment,
targetAssignment,
+ currentPartitionAssignment, targetPartitionAssignment)) {
+ // Check if the servers of the first assignment for the group have any
space left to move the segments of the
+ // group as a whole. If so, let's mark the group as to be moved,
otherwise we mark the group as a whole as not
+ // moveable. It is enough to check the first segment (or partition)
since all segments and partitions in the
+ // group have the same current and target instances.
+ Map<String, String> firstCurrentInstanceStateMap;
+ Map<String, String> firstTargetInstanceStateMap;
+ if (!group._segments.isEmpty()) {
+ Map.Entry<String, Map<String, String>> firstEntry =
group._segments.entrySet().iterator().next();
+ firstCurrentInstanceStateMap = firstEntry.getValue();
+ firstTargetInstanceStateMap =
targetAssignment.get(firstEntry.getKey());
+ } else {
+ int firstPartitionId = group._partitionIds.get(0);
+ firstCurrentInstanceStateMap =
getPartitionInstanceStateMap(currentPartitionAssignment.get(firstPartitionId));
+ firstTargetInstanceStateMap =
getPartitionInstanceStateMap(targetPartitionAssignment.get(firstPartitionId));
+ }
+ SingleSegmentAssignment firstAssignment =
+ getNextSingleSegmentAssignment(firstCurrentInstanceStateMap,
firstTargetInstanceStateMap,
+ minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap);
+ Set<String> serversAdded =
+
getServersAddedInSingleSegmentAssignment(firstCurrentInstanceStateMap,
firstAssignment._instanceStateMap);
+ int numSegments = group._segments.size();
+ boolean anyServerExhaustedBatchSize = false;
+ for (String server : serversAdded) {
+ int segmentsAddedToServerSoFar =
serverToNumSegmentsAddedSoFar.getOrDefault(server, 0);
+ // Case I: We already exceeded the batchSizePerServer for this server,
cannot add any more segments
+ // Case II: We have not yet exceeded the batchSizePerServer for this
server, but we don't have sufficient
+ // space to host the segments for this assignment on the server, and
we have allocated some partitions so
+ // far. If the batchSizePerServer is less than the number of segments
in a given partitionId, we must host
+ // at least 1 partition and exceed the batchSizePerServer to ensure
progress is made. Thus, performing this
+ // check only if segmentsAddedToServerSoFar > 0 is necessary.
+ if ((segmentsAddedToServerSoFar >= batchSizePerServer) ||
(segmentsAddedToServerSoFar > 0
+ && (segmentsAddedToServerSoFar + numSegments) >
batchSizePerServer)) {
+ anyServerExhaustedBatchSize = true;
+ break;
+ }
+ }
+ moveStrictReplicaGroup(group, anyServerExhaustedBatchSize,
targetAssignment, currentPartitionAssignment,
+ targetPartitionAssignment, nextAssignment, nextPartitionAssignment,
minAvailableReplicas, lowDiskMode,
+ numSegmentsToOffloadMap, assignmentMap, availableInstancesMap,
serverToNumSegmentsAddedSoFar,
+ dataLossRiskAssessor);
}
checkIfAnyServersAssignedMoreSegmentsThanBatchSize(batchSizePerServer,
serverToNumSegmentsAddedSoFar,
tableRebalanceLogger);
- return nextAssignment;
+ return new TableAssignment(nextAssignment, nextPartitionAssignment);
}
- private static void updateNextAssignmentForPartitionIdStrictReplicaGroup(
- Map<String, Map<String, String>> currentAssignment, Map<String,
Map<String, String>> targetAssignment,
- Map<String, Map<String, String>> nextAssignment, boolean
anyServerExhaustedBatchSize, int minAvailableReplicas,
- boolean lowDiskMode, Map<String, Integer> numSegmentsToOffloadMap,
- Map<Pair<Set<String>, Set<String>>, Set<String>> assignmentMap,
+ /// Moves the segments and the partitions of the given group to their next
assignment, or keeps them as they are
+ /// when the batch size is exhausted. The segments and partitions of a group
have the same current and target
+ /// instances, so they take the same step; the minimum available replicas
check is still performed per segment and
+ /// partition because it depends on the other segments assigned to the same
instances.
+ private static void moveStrictReplicaGroup(StrictReplicaGroup group, boolean
anyServerExhaustedBatchSize,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, Map<String,
Map<String, String>> nextAssignment,
+ Map<Integer, List<String>> nextPartitionAssignment, int
minAvailableReplicas, boolean lowDiskMode,
+ Map<String, Integer> numSegmentsToOffloadMap, Map<Pair<Set<String>,
Set<String>>, Set<String>> assignmentMap,
Map<Set<String>, Set<String>> availableInstancesMap, Map<String,
Integer> serverToNumSegmentsAddedSoFar,
DataLossRiskAssessor dataLossRiskAssessor) {
if (anyServerExhaustedBatchSize) {
- // Exhausted the batch size for at least 1 server, just copy over the
remaining segments as is
- nextAssignment.putAll(currentAssignment);
- } else {
- // Process all the partitionIds even if segmentsAddedSoFar becomes
larger than batchSizePerServer
- // Can only do bestEfforts w.r.t. StrictReplicaGroup since a whole
partition must be moved together for
- // maintaining consistency
- for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
- String segmentName = entry.getKey();
- Map<String, String> currentInstanceStateMap = entry.getValue();
- Map<String, String> targetInstanceStateMap =
targetAssignment.get(segmentName);
- SingleSegmentAssignment assignment =
- getNextSingleSegmentAssignment(currentInstanceStateMap,
targetInstanceStateMap, minAvailableReplicas,
- lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
- Set<String> assignedInstances = assignment._instanceStateMap.keySet();
- Set<String> availableInstances = assignment._availableInstances;
- availableInstancesMap.compute(assignedInstances, (k,
currentAvailableInstances) -> {
- if (currentAvailableInstances == null) {
- // First segment assigned to these instances, use the new
assignment and update the available instances
- nextAssignment.put(segmentName, assignment._instanceStateMap);
- updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(), k);
- return availableInstances;
- } else {
- // There are other segments assigned to the same instances, check
the available instances to see if
- // adding the new assignment can still hold the minimum available
replicas requirement
- availableInstances.retainAll(currentAvailableInstances);
- if (availableInstances.size() >= minAvailableReplicas) {
- // New assignment can be added
- nextAssignment.put(segmentName, assignment._instanceStateMap);
- updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(), k);
- return availableInstances;
- } else {
- // New assignment cannot be added, use the current instance
state map
- nextAssignment.put(segmentName, currentInstanceStateMap);
- return currentAvailableInstances;
- }
- }
- });
+ // Exhausted the batch size for at least 1 server, just copy over the
remaining segments and partitions as is
+ nextAssignment.putAll(group._segments);
+ for (int partitionId : group._partitionIds) {
+ nextPartitionAssignment.put(partitionId,
currentPartitionAssignment.get(partitionId));
+ }
+ return;
+ }
+ // Process all the segments even if segmentsAddedSoFar becomes larger than
batchSizePerServer
+ // Can only do bestEfforts w.r.t. StrictReplicaGroup since a whole
partition must be moved together for
+ // maintaining consistency
+ for (Map.Entry<String, Map<String, String>> entry :
group._segments.entrySet()) {
+ String segmentName = entry.getKey();
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ Map<String, String> targetInstanceStateMap =
targetAssignment.get(segmentName);
+ SingleSegmentAssignment assignment =
+ getNextSingleSegmentAssignment(currentInstanceStateMap,
targetInstanceStateMap, minAvailableReplicas,
+ lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
+ Map<String, String> nextInstanceStateMap;
+ if (canAssignWithStrictReplicaGroup(assignment, minAvailableReplicas,
availableInstancesMap)) {
+ nextInstanceStateMap = assignment._instanceStateMap;
+ updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(),
+ nextInstanceStateMap.keySet());
+ } else {
+ nextInstanceStateMap = currentInstanceStateMap;
+ }
+ nextAssignment.put(segmentName, nextInstanceStateMap);
- if (!nextAssignment.get(segmentName).equals(currentInstanceStateMap)) {
- Set<String> serversAddedForSegment =
-
getServersAddedInSingleSegmentAssignment(currentInstanceStateMap,
nextAssignment.get(segmentName));
- serversAddedForSegment.forEach(server ->
serverToNumSegmentsAddedSoFar.merge(server, 1, Integer::sum));
+ if (!nextInstanceStateMap.equals(currentInstanceStateMap)) {
+ Set<String> serversAddedForSegment =
+ getServersAddedInSingleSegmentAssignment(currentInstanceStateMap,
nextInstanceStateMap);
+ serversAddedForSegment.forEach(server ->
serverToNumSegmentsAddedSoFar.merge(server, 1, Integer::sum));
- // Since next assignment doesn't match current assignment, it means
the segment will be moved. Check if there
- // is a data loss risk
- Pair<Boolean, String> dataLossRiskResult =
dataLossRiskAssessor.assessDataLossRisk(segmentName);
- if (dataLossRiskResult.getLeft()) {
- throw new IllegalStateException(dataLossRiskResult.getRight());
- }
+ // Since next assignment doesn't match current assignment, it means
the segment will be moved. Check if there
+ // is a data loss risk
+ Pair<Boolean, String> dataLossRiskResult =
dataLossRiskAssessor.assessDataLossRisk(segmentName);
+ if (dataLossRiskResult.getLeft()) {
+ throw new IllegalStateException(dataLossRiskResult.getRight());
}
}
}
+ for (int partitionId : group._partitionIds) {
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ SingleSegmentAssignment assignment =
+
getNextSingleSegmentAssignment(getPartitionInstanceStateMap(currentInstances),
+
getPartitionInstanceStateMap(targetPartitionAssignment.get(partitionId)),
minAvailableReplicas,
+ lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
+ nextPartitionAssignment.put(partitionId,
+ canAssignWithStrictReplicaGroup(assignment, minAvailableReplicas,
availableInstancesMap)
Review Comment:
Fixed in 7df9c15: a partition entry follows the step decided for the
segments with its current and target instances, and an entry without such
segments is stepped without the minimum available replicas check, so entries no
longer register available instances.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -21,47 +21,97 @@
import com.google.common.base.Preconditions;
import java.util.List;
import java.util.Map;
+import java.util.TreeMap;
import javax.annotation.Nullable;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.helix.model.IdealState;
import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.common.assignment.InstancePartitionsUtils;
+import org.apache.pinot.common.assignment.PartitionAssignmentUtils;
import org.apache.pinot.common.restlet.resources.RebalanceConfig;
import org.apache.pinot.common.tier.Tier;
+import org.apache.pinot.common.utils.SegmentUtils;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.AllServersSegmentAssignmentStrategy;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategy;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategyFactory;
+import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
/// Segment assignment for offline table.
+///
+/// For a partition-pinned table (strict replica-group routing), all segments
of a value partition are assigned to
+/// the same instances, recorded in the partition assignment of the ideal
state (see
+/// [org.apache.pinot.common.assignment.PartitionAssignmentUtils]), even when
the instance partitions have changed.
+/// The partition id of a new segment is read from its ZK metadata. This
requires a partition column, a table that is
+/// not a dimension table (whose segments are assigned to all servers of the
tenant), and OFFLINE instance partitions
+/// holding exactly one instance per partition within each replica-group;
otherwise a warning is logged and the
+/// segments are assigned with the segment assignment strategy, as are
segments without partition id and segments
+/// eligible for a tier.
public class OfflineSegmentAssignment extends BaseSegmentAssignment {
@Override
- public List<String> assignSegment(String segmentName, Map<String,
Map<String, String>> currentAssignment,
+ protected boolean isPartitionPinned(TableConfig tableConfig) {
+ if (!super.isPartitionPinned(tableConfig)) {
+ return false;
+ }
+ if (_partitionColumn == null) {
+ _logger.warn("Table: {} with strict replica-group routing has no
partition column, cannot pin its segments to "
+ + "their partition, assigning them with the segment assignment
strategy", _tableNameWithType);
+ return false;
+ }
+ if (tableConfig.isDimTable()) {
+ _logger.warn("Dimension table: {} with strict replica-group routing
cannot pin its segments to their partition "
+ + "because they are assigned to all servers, assigning them with the
segment assignment strategy",
+ _tableNameWithType);
+ return false;
+ }
+ return true;
+ }
+
+ @Override
+ public List<String> assignSegment(String segmentName, IdealState idealState,
Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
InstancePartitions instancePartitions =
instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
Preconditions.checkState(instancePartitions != null, "Failed to find
OFFLINE instance partitions for table: %s",
_tableNameWithType);
- // Gets Segment assignment strategy for instance partitions
- SegmentAssignmentStrategy segmentAssignmentStrategy =
-
SegmentAssignmentStrategyFactory.getSegmentAssignmentStrategy(_helixManager,
_tableConfig,
- InstancePartitionsType.OFFLINE.toString(), instancePartitions);
_logger.info("Assigning segment: {} with instance partitions: {} for
table: {}", segmentName, instancePartitions,
_tableNameWithType);
- List<String> instancesAssigned =
- segmentAssignmentStrategy.assignSegment(segmentName,
currentAssignment, instancePartitions);
+ List<String> instancesAssigned = null;
+ if (isPinnedWith(instancePartitions)) {
+ Integer partitionId =
+ SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType,
_helixManager, _partitionColumn);
+ if (partitionId != null) {
+ instancesAssigned = getOrCreatePartitionInstances(idealState,
partitionId,
Review Comment:
Documented in the class Javadoc: the first entry of a partition of a
pre-existing offline table is created from the instance partitions, where the
strategy would place the segment. An offline table has no canonical segment to
seed from, and no collocation requirement for correctness; the rebalance
rewrites the entries.
##########
pinot-common/src/main/java/org/apache/pinot/common/assignment/PartitionAssignmentUtils.java:
##########
@@ -0,0 +1,102 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.assignment;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import javax.annotation.Nullable;
+import org.apache.helix.model.IdealState;
+
+
+/// Partition assignment of a table: the instances hosting all segments of
each segment partition, stored in the list
+/// fields of the table IdealState. List fields whose key is not a partition
id are not part of it and are left alone.
+///
+/// Each list field maps a segment partition id (the stream partition for a
real-time table, the value partition for an
+/// offline table) to the sorted list of instances hosting all segments of
that partition. The partition assignment is
+/// maintained for partition-pinned tables only, i.e. tables whose segments
must be collocated per partition: tables
+/// with strict replica-group routing, which upsert and dedup tables use. For
these tables:
+/// - A new segment is assigned to the instances of its partition when the
partition has an entry, even when the
+/// instance partitions have changed since the entry was created. The entry
is created from the instance partitions
+/// otherwise.
+/// - Table rebalance moves the entries together with the segments, so that a
segment added during the rebalance is
+/// assigned to the same instances as the other segments of its partition at
that point of the rebalance.
+///
+/// The entries are read and written together with the map fields of the
IdealState, so they are never observed out of
+/// sync with the segment assignment. Helix ignores the list fields of an
IdealState in CUSTOMIZED rebalance mode.
+public class PartitionAssignmentUtils {
+ private PartitionAssignmentUtils() {
+ }
+
+ /// Returns the instances of the given partition, or `null` if the partition
has no entry.
+ @Nullable
+ public static List<String> getInstances(IdealState idealState, int
partitionId) {
+ return idealState.getRecord().getListField(Integer.toString(partitionId));
+ }
+
+ /// Sets the instances of the given partition.
+ public static void setInstances(IdealState idealState, int partitionId,
Collection<String> instances) {
+ idealState.getRecord().setListField(Integer.toString(partitionId),
toInstanceList(instances));
+ }
+
+ /// Returns the partition assignment as a sorted map from partition id to
the instances of the partition. The lists
+ /// are the ones stored in the IdealState record and must not be modified.
+ public static Map<Integer, List<String>> getPartitionAssignment(IdealState
idealState) {
+ Map<Integer, List<String>> partitionAssignment = new TreeMap<>();
+ for (Map.Entry<String, List<String>> entry :
idealState.getRecord().getListFields().entrySet()) {
+ try {
+ partitionAssignment.put(Integer.parseInt(entry.getKey()),
entry.getValue());
Review Comment:
Kept as is: the list fields of a Pinot IdealState are only written by this
code, with `Integer.toString(id)` keys, so the parse stays lenient.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
