Jackie-Jiang commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4342832710
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -396,6 +402,30 @@ private RebalanceResult doRebalance(TableConfig
tableConfig, RebalanceConfig reb
if (segmentAssignmentUnchanged) {
tableRebalanceLogger.info("Table is already balanced");
+ if (!partitionAssignmentUnchanged && !dryRun) {
+ // Only the partition assignment changed (e.g. a table created before
the partition assignment was introduced,
+ // or partitions without segment): update it without moving any segment
+ tableRebalanceLogger.info("Updating the partition assignment from: {}
to: {}", currentPartitionAssignment,
+ targetPartitionAssignment);
+ PartitionAssignmentUtils.setPartitionAssignment(currentIdealState,
targetPartitionAssignment);
+ ZNRecord idealStateRecord = currentIdealState.getRecord();
+ try {
+ Preconditions.checkState(_helixDataAccessor.getBaseDataAccessor()
+ .set(idealStatePropertyKey.getPath(), idealStateRecord,
idealStateRecord.getVersion(),
Review Comment:
Fixed in 7df9c15: the partition-only update goes through the rebalance loop,
with its version check and retries, and a dry run returns DONE for such a table
so the REST flow proceeds to the actual rebalance.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -158,8 +178,100 @@ protected List<String> assignConsumingSegment(int
segmentPartitionId, InstancePa
}
}
+ /// Assigns a CONSUMING segment of a partition-pinned table to the instances
of its partition, see
+ /// [BaseSegmentAssignment#getOrCreatePartitionInstances].
+ private List<String> assignPinnedConsumingSegment(String segmentName,
IdealState idealState,
+ InstancePartitions instancePartitions) {
+ Integer partitionId =
+ SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType,
_helixManager, _partitionColumn);
+ Preconditions.checkState(partitionId != null, "Failed to find partition id
for segment: %s of table: %s",
+ segmentName, _tableNameWithType);
+ return getOrCreatePartitionInstances(idealState, partitionId,
Review Comment:
Kept as is and documented in the description: a pinned table needs a
rebalance after its instance partitions change, as upsert/dedup tables do
today. Letting the candidate win when the entry's instances are not in the
instance partitions would also fire mid-rebalance, where the entry steps away
from the old instances.
##########
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) {
Review Comment:
Fixed in 7df9c15: the fast path only re-targets new segments when the
entries are stepped (real-time table with consuming segments included). For
offline tables it keeps new segments as they are, and they already land on the
target entries. No recomputation per upload, and a tier-eligible upload stays
where the tiered assignment put it.
##########
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()
Review Comment:
Fixed in 7df9c15: a segment added to a partition without entry triggers a
full recalculation of the target. Covered in `TableRebalancerTest`.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -158,8 +178,100 @@ protected List<String> assignConsumingSegment(int
segmentPartitionId, InstancePa
}
}
+ /// Assigns a CONSUMING segment of a partition-pinned table to the instances
of its partition, see
+ /// [BaseSegmentAssignment#getOrCreatePartitionInstances].
+ private List<String> assignPinnedConsumingSegment(String segmentName,
IdealState idealState,
+ InstancePartitions instancePartitions) {
+ Integer partitionId =
+ SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType,
_helixManager, _partitionColumn);
+ Preconditions.checkState(partitionId != null, "Failed to find partition id
for segment: %s of table: %s",
+ segmentName, _tableNameWithType);
+ return getOrCreatePartitionInstances(idealState, partitionId,
+ assignConsumingSegment(partitionId, instancePartitions));
+ }
+
+ /// OFFLINE segments are skipped because they are not rebalanced, so their
assignment can be stale. The instances are
+ /// taken from, in order:
+ /// - The CONSUMING segment of the partition, which is always hosted by the
CONSUMING instances of the partition.
+ /// - The latest LLC segment of the partition, else an uploaded segment of
the partition, whose partition id is in
+ /// its name or in its ZK metadata. Completed segments are not used when
they are relocated to the COMPLETED
+ /// instances, or when the segment has been moved to a tier, because they
are not on the CONSUMING instances of
+ /// the partition then. The tier is read from the segment ZK metadata
whether or not the table still has tiers
+ /// configured, because removing the tier config does not move the
segments back.
+ @Nullable
@Override
- public Map<String, Map<String, String>> rebalanceTable(Map<String,
Map<String, String>> currentAssignment,
+ protected Set<String> getExistingAssignment(int partitionId, Map<String,
Map<String, String>> currentAssignment) {
+ LLCSegmentName latestLLCSegmentName = null;
+ Map<String, String> latestInstanceStateMap = null;
+ List<String> uploadedSegments = new ArrayList<>();
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segmentName = entry.getKey();
+ Map<String, String> instanceStateMap = entry.getValue();
+ boolean completed;
+ if (instanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
+ completed = true;
+ } else if (instanceStateMap.containsValue(SegmentStateModel.CONSUMING)) {
+ completed = false;
+ } else {
+ // All instances are OFFLINE, and OFFLINE segments are not rebalanced,
so their assignment can be stale
+ continue;
+ }
+ LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
+ if (llcSegmentName == null) {
+ uploadedSegments.add(segmentName);
+ continue;
+ }
+ if (llcSegmentName.getPartitionGroupId() != partitionId) {
+ continue;
+ }
+ if (!completed) {
+ return instanceStateMap.keySet();
+ }
+ if (latestLLCSegmentName == null
+ || llcSegmentName.getSequenceNumber() >
latestLLCSegmentName.getSequenceNumber()) {
+ latestLLCSegmentName = llcSegmentName;
+ latestInstanceStateMap = instanceStateMap;
+ }
+ }
+ // Completed segments relocated to the COMPLETED instances are not on the
CONSUMING instances of the partition
+ if
(InstanceAssignmentConfigUtils.shouldRelocateCompletedSegments(_tableConfig)) {
+ return null;
+ }
+ String latestSegmentName = null;
+ if (latestLLCSegmentName != null) {
+ latestSegmentName = latestLLCSegmentName.getSegmentName();
+ } else {
+ for (String uploadedSegment : uploadedSegments) {
+ Integer uploadedSegmentPartitionId =
+ SegmentUtils.getSegmentPartitionId(uploadedSegment,
_tableNameWithType, _helixManager, _partitionColumn);
Review Comment:
Fixed in 7df9c15: the relocation check comes first, uploaded segments are
resolved with the bulk read shared by the assignment instance, and a segment
without ZK metadata is skipped. Covered in
`PartitionPinnedRealtimeSegmentAssignmentTest`.
--
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]
