Re: [PR] Skip full target recompute on IdealState version change for strict realtime rebalance without instance reassignment [pinot]

["Jackie-Jiang (via GitHub)" <gi...@apache.org>]

Jackie-Jiang commented on code in PR #19054:
URL: https://github.com/apache/pinot/pull/19054#discussion_r3636117612


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -2233,20 +2248,40 @@ static Set<String> getMovingConsumingSegments(Map<String, Map<String, String>> c
       Map<String, Map<String, String>> targetAssignment) {
     Set<String> movingConsumingSegments = new HashSet<>();
     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);
-      if (targetInstanceStateMap != null && targetInstanceStateMap.values().stream()
-          .noneMatch(state -> state.equals(SegmentStateModel.ONLINE)) && targetInstanceStateMap.values().stream()
-          .anyMatch(state -> state.equals(SegmentStateModel.CONSUMING))) {
-        if (!currentInstanceStateMap.keySet().equals(targetInstanceStateMap.keySet())) {
-          movingConsumingSegments.add(segmentName);
-        }
+      if (isMovingConsumingSegment(entry.getValue(), targetAssignment.get(entry.getKey()))) {
+        movingConsumingSegments.add(entry.getKey());
       }
     }
     return movingConsumingSegments;
   }
 
+  /// Returns whether any consuming segment moves between the current and target assignment, short-circuiting on the
+  /// first such segment. Cheaper than `!getMovingConsumingSegments(...).isEmpty()` when only existence is needed
+  /// because it does not materialize the full set.
+  @VisibleForTesting
+  static boolean hasMovingConsumingSegments(Map<String, Map<String, String>> currentAssignment,
+      Map<String, Map<String, String>> targetAssignment) {
+    for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) {
+      if (isMovingConsumingSegment(entry.getValue(), targetAssignment.get(entry.getKey()))) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /// Returns whether the segment is a consuming segment being moved: its target state is `CONSUMING` and its assigned
+  /// instances differ between the current and target assignment. A segment's target instance state map is uniform
+  /// (every replica is `ONLINE`, `CONSUMING`, or `OFFLINE`) because `RealtimeSegmentAssignment` assigns a single state
+  /// to all instances of a moved segment and keeps `OFFLINE` segments unchanged, so the segment's state can be read
+  /// from the first entry. The map is non-null and non-empty: callers iterate the current assignment and look up a
+  /// target derived from it (via `rebalanceTable` / `getNextAssignment`), which covers every current segment, and
+  /// every segment has at least one replica.
+  private static boolean isMovingConsumingSegment(Map<String, String> currentInstanceStateMap,
+      Map<String, String> targetInstanceStateMap) {
+    return targetInstanceStateMap.values().iterator().next().equals(SegmentStateModel.CONSUMING)

Review Comment:
   Good catch, and thanks for the detailed scenario. The PR has been reworked and no longer touches `getMovingConsumingSegments()` at all — the first-entry inference and the removed mixed-state test case are both restored to the original code.
   
   The recompute gate is now based on the rebalance config instead: for strict realtime assignment, the full target recompute on an IdealState version change is skipped only when the rebalance is not reassigning instances (`reassignInstances`/`bootstrap` both false, e.g. the `SegmentRelocator` tier-relocation flow). In that case the instance partitions are read-only for the job, so a segment added mid-rebalance is placed consistently with the target. With instance reassignment, the conservative full recompute is kept, matching prior behavior.
   



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -2233,20 +2248,40 @@ static Set<String> getMovingConsumingSegments(Map<String, Map<String, String>> c
       Map<String, Map<String, String>> targetAssignment) {
     Set<String> movingConsumingSegments = new HashSet<>();
     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);
-      if (targetInstanceStateMap != null && targetInstanceStateMap.values().stream()
-          .noneMatch(state -> state.equals(SegmentStateModel.ONLINE)) && targetInstanceStateMap.values().stream()
-          .anyMatch(state -> state.equals(SegmentStateModel.CONSUMING))) {
-        if (!currentInstanceStateMap.keySet().equals(targetInstanceStateMap.keySet())) {
-          movingConsumingSegments.add(segmentName);
-        }
+      if (isMovingConsumingSegment(entry.getValue(), targetAssignment.get(entry.getKey()))) {
+        movingConsumingSegments.add(entry.getKey());
       }
     }
     return movingConsumingSegments;
   }
 
+  /// Returns whether any consuming segment moves between the current and target assignment, short-circuiting on the
+  /// first such segment. Cheaper than `!getMovingConsumingSegments(...).isEmpty()` when only existence is needed
+  /// because it does not materialize the full set.
+  @VisibleForTesting
+  static boolean hasMovingConsumingSegments(Map<String, Map<String, String>> currentAssignment,
+      Map<String, Map<String, String>> targetAssignment) {
+    for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) {
+      if (isMovingConsumingSegment(entry.getValue(), targetAssignment.get(entry.getKey()))) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /// Returns whether the segment is a consuming segment being moved: its target state is `CONSUMING` and its assigned
+  /// instances differ between the current and target assignment. A segment's target instance state map is uniform
+  /// (every replica is `ONLINE`, `CONSUMING`, or `OFFLINE`) because `RealtimeSegmentAssignment` assigns a single state
+  /// to all instances of a moved segment and keeps `OFFLINE` segments unchanged, so the segment's state can be read
+  /// from the first entry. The map is non-null and non-empty: callers iterate the current assignment and look up a
+  /// target derived from it (via `rebalanceTable` / `getNextAssignment`), which covers every current segment, and
+  /// every segment has at least one replica.
+  private static boolean isMovingConsumingSegment(Map<String, String> currentInstanceStateMap,
+      Map<String, String> targetInstanceStateMap) {
+    return targetInstanceStateMap.values().iterator().next().equals(SegmentStateModel.CONSUMING)

Review Comment:
   The PR has been reworked and no longer modifies `getMovingConsumingSegments()` — the original `noneMatch(ONLINE)`/`anyMatch(CONSUMING)` classification is untouched. The recompute gate is now based on `reassignInstances`/`bootstrap` instead of consuming-segment movement.
   



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -585,6 +585,14 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb
         new PartitionIdFetcherImpl(tableNameWithType, TableConfigUtils.getPartitionColumn(tableConfig), _helixManager,
             isStrictRealtimeSegmentAssignment);
 
+    // For strict realtime segment assignment, whether this rebalance moves any consuming segment (e.g. rebalancing the
+    // consuming tier) as opposed to only completed segments (e.g. a tier relocation). This is the only case where a

Review Comment:
   You're right — and your paused-table scenario in the review comment showed the consuming-segment-based gate was insufficient in general (segments uploaded by a minion task mid-rebalance would not be re-collocated). The gate is now based on instance reassignment instead: a newly added segment is placed from the current instance partitions, so as long as the rebalance is not reassigning instances (`reassignInstances`/`bootstrap` both false), it is placed consistently with the target and the cheap path is safe — regardless of whether consuming segments move. With reassignment enabled, the full recompute is kept as before. The comment this nit targeted has been rewritten accordingly.
   



-- 
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@pinot.apache.org

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


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