You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by "Jackie-Jiang (via GitHub)" <gi...@apache.org> on 2026/10/17 02:08:30 UTC

[PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Jackie-Jiang opened a new pull request, #19916:
URL: https://github.com/apache/pinot/pull/19916

   ## Summary
   
   Strict replica-group routing requires all segments of a partition to be hosted by the same instances. Today this is only enforced for upsert and dedup real-time tables, through `BaseStrictRealtimeSegmentAssignment`, which on every new segment scans the IdealState for an existing segment of the partition (reading segment ZK metadata for uploaded segments, and the latest segment's tier for dedup tables), and `TableRebalancer` reads partition ids from segment ZK metadata again to batch segments per partition. Offline tables with strict replica-group routing get no such guarantee.
   
   This PR records the **partition assignment** in the list fields of the table IdealState: one list field per segment partition id, holding the sorted instances hosting all segments of that partition. It is maintained for tables with strict replica-group routing (which upsert and dedup tables use):
   - Real-time tables: the entries describe the CONSUMING placement. CONSUMING segments, and COMPLETED segments not relocated to COMPLETED instance partitions or a tier, follow them.
   - Offline tables with a partition column whose OFFLINE instance partitions hold one instance per partition within each replica-group: non-tier segments follow them.
   
   The entries are written in the same IdealState update as the map fields, so they are never out of sync with the segment assignment. Helix ignores IdealState list fields in CUSTOMIZED mode.
   
   ### Segment assignment
   - A new segment is assigned to the entry of its partition in O(1), without scanning the IdealState or reading other segments' ZK metadata. The partition id comes from the segment name, or from the segment's own ZK metadata for offline and custom-named uploaded segments (the same single read the replica-group strategy does today).
   - When the partition has no entry, the entry is created from the existing segments of the partition (real-time tables created before this change, see below), else from the instance partitions. The existing mismatch meter still fires when an entry differs from what the instance partitions give.
   - `SegmentAssignment.assignSegment` takes the `IdealState` instead of its map fields so it can read and write the entries. `rebalanceTable` takes the `IdealState` and returns a `TableAssignment` (segment assignment + partition assignment).
   - `BaseStrictRealtimeSegmentAssignment`, `SingleTierStrictRealtimeSegmentAssignment` and `MultiTierStrictRealtimeSegmentAssignment` are removed: `RealtimeSegmentAssignment` and `OfflineSegmentAssignment` pin the segments through shared helpers in `BaseSegmentAssignment`. The multi-tier special case of checking the tier of the latest segment is gone, because the entry is the CONSUMING placement regardless of where old segments went.
   
   ### Rebalance
   - `rebalanceTable` computes the target instances of every partition from the instance partitions, and the segments follow their partition. Offline tables resolve the partition ids of all segments with a single bulk ZK read per rebalance instead of one read per segment.
   - The rebalance loop moves the entries along with the segments at every IdealState write: a partition takes the same step as the segments with the same current and target instances, so a segment added mid-rebalance (assigned to the entry) lands on the same instances as its siblings and is then moved with them. In strict replica-group mode the batching unit is the group of segments with the same current and target instances, sub-grouped by the partition id in the segment names; a group with a segment lacking a partition id in its name moves as a whole. This removes `PartitionIdFetcher` and its ZK reads.
   - Downtime rebalance writes the target entries with the target assignment. Convergence checks the entries as well as the segments.
   - Tables created before this change have no entries. Their first rebalance populates them: alone when no segment moves, or once the segments reach the target otherwise. Meanwhile a new real-time segment of a partition without entry is placed by the existing segments of the partition, as today.
   
   ### Behavior changes
   - Offline tables with strict replica-group routing, a partition column and one instance per partition within each replica-group now keep a new segment with the other segments of its partition when the instance partitions have changed, instead of placing it by the new instance partitions.
   - Real-time tables with strict replica-group routing that are not upsert or dedup tables get the same pinned assignment.
   - Upsert and dedup tables configured with the beta `multiStageReplicaGroup` routing are no longer pinned: pinning is decided by the routing config alone. Their validation is unchanged.
   - An old controller ignores the entries and does not update them on rebalance, so controllers should be upgraded before rebalancing; the first rebalance by a new controller rewrites all entries.
   
   The broker still ignores the list fields; using them in `SegmentPartitionMetadataManager` is a follow-up.
   


-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang merged PR #19916:
URL: https://github.com/apache/pinot/pull/19916


-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
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: 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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #19916:
URL: https://github.com/apache/pinot/pull/19916#issuecomment-6228873111

   ## [Codecov](https://app.codecov.io/gh/apache/pinot/pull/19916?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
   :white_check_mark: All modified and coverable lines are covered by tests.
   :white_check_mark: Project coverage is 58.95%. Comparing base ([`a45ebf0`](https://app.codecov.io/gh/apache/pinot/commit/a45ebf01379257591c27c900c3525612a3c49b92?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`d7ec1b6`](https://app.codecov.io/gh/apache/pinot/commit/d7ec1b61d3dafb898ffdcca9956e86d0af5b7bf1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
   > :exclamation:  There is a different number of reports uploaded between BASE (a45ebf0) and HEAD (d7ec1b6). Click for more details.
   > 
   > <details><summary>HEAD has 24 uploads less than BASE</summary>
   >
   >| Flag | BASE (a45ebf0) | HEAD (d7ec1b6) |
   >|------|------|------|
   >|lane-a|2|0|
   >|integration|4|0|
   >|java-25|6|1|
   >|integration1|2|0|
   >|temurin|6|1|
   >|lane-b|2|0|
   >|unittests|2|1|
   >|unittests2|1|0|
   >|integration2|2|0|
   ></details>
   
   <details><summary>Additional details and impacted files</summary>
   
   
   
   ```diff
   @@              Coverage Diff              @@
   ##             master   #19916       +/-   ##
   =============================================
   - Coverage     69.49%   58.95%   -10.55%     
   + Complexity     1491        1     -1490     
   =============================================
     Files          3546     2730      -816     
     Lines        234166   171764    -62402     
     Branches      37493    28307     -9186     
   =============================================
   - Hits         162739   101266    -61473     
   - Misses        58854    61954     +3100     
   + Partials      12573     8544     -4029     
   ```
   
   | [Flag](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
   |---|---|---|
   | [integration](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [integration1](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [integration2](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [java-25](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `58.95% <100.00%> (-10.55%)` | :arrow_down: |
   | [lane-a](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [lane-b](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [temurin](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `58.95% <100.00%> (-10.55%)` | :arrow_down: |
   | [unittests](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `58.95% <100.00%> (-10.54%)` | :arrow_down: |
   | [unittests1](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `58.95% <100.00%> (+0.03%)` | :arrow_up: |
   | [unittests2](https://app.codecov.io/gh/apache/pinot/pull/19916/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
   </details>
   
   [:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/pinot/pull/19916?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
   :loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   <details><summary> :rocket: New features to boost your workflow: </summary>
   
   - :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
   </details>


-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4342830085


##########
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",

Review Comment:
   Fixed in 7df9c15: a segment without a derivable partition id now falls back to the non-pinned assignment (partition id derived from the name) with a warning, both on assignment and on rebalance. Upsert and dedup tables still reject it, as before this PR. Covered in `PartitionPinnedRealtimeSegmentAssignmentTest`.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   @Nullable
   private Map<String, String> _segmentToTierName;
 
+  // Cache of segment name to partition id (null value means the partition id cannot be determined) for the segments
+  // whose partition id is not in the segment name, loaded with a single bulk ZK read on first use. Same scope and
+  // limitations as _segmentToTierName; a segment absent from the cache is resolved on demand and merged in.
+  @Nullable
+  private Map<String, Integer> _segmentToPartitionId;
+
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig, @Nullable ControllerMetrics controllerMetrics) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
     _tableConfig = tableConfig;
     _replication = tableConfig.getReplication();
     _partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
-    _controllerMetrics = controllerMetrics;
+    _partitionPinned = isPartitionPinned(tableConfig);
     if (_partitionColumn == null) {
-      _logger.info("Initialized with replication: {} without partition column for table: {} ", _replication,
-          _tableNameWithType);
+      _logger.info("Initialized with replication: {} without partition column, partitionPinned: {} for table: {}",
+          _replication, _partitionPinned, _tableNameWithType);
     } else {
-      _logger.info("Initialized with replication: {} and partition column: {} for table: {}", _replication,
-          _partitionColumn, _tableNameWithType);
+      _logger.info("Initialized with replication: {} and partition column: {}, partitionPinned: {} for table: {}",
+          _replication, _partitionColumn, _partitionPinned, _tableNameWithType);
+    }
+  }
+
+  /// Returns whether the segments of the table are pinned to their partition (see the class Javadoc): the table uses
+  /// strict replica-group routing. Subclasses add the conditions they need to pin the segments, and log a warning
+  /// when such a table cannot be pinned, in which case its segments are assigned with the segment assignment strategy.
+  protected boolean isPartitionPinned(TableConfig tableConfig) {
+    return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);
+  }
+
+  /// Returns the instances of the given partition of a partition-pinned table, creating its entry in the partition
+  /// assignment of the given ideal state when the partition has no entry yet, or an empty one (a rebalance in low disk
+  /// mode without minimum available replicas empties the entry while it drops the partition from its current
+  /// instances). The entry is created from the existing assignment of the partition (see [#getExistingAssignment])
+  /// when there is one, else from the given candidate instances computed from the instance partitions. An existing
+  /// entry wins over the candidate instances, so that the segments of the partition stay collocated when the instance
+  /// partitions change: the partition moves on the next rebalance.
+  protected List<String> getOrCreatePartitionInstances(IdealState idealState, int partitionId,
+      List<String> candidateInstances) {
+    List<String> instances = PartitionAssignmentUtils.getInstances(idealState, partitionId);
+    if (CollectionUtils.isEmpty(instances)) {

Review Comment:
   Kept as is and documented in the description: an old controller does not maintain the entries, so a rebalance run on the old version (rolling upgrade or rollback) leaves them stale until the first rebalance by the new version. Validating an entry against the partition's CONSUMING segment would bring back the IdealState scan on every new segment, and disagrees with rebalances that exclude consuming segments, where the entry legitimately differs from the CONSUMING segment.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   @Nullable
   private Map<String, String> _segmentToTierName;
 
+  // Cache of segment name to partition id (null value means the partition id cannot be determined) for the segments
+  // whose partition id is not in the segment name, loaded with a single bulk ZK read on first use. Same scope and
+  // limitations as _segmentToTierName; a segment absent from the cache is resolved on demand and merged in.
+  @Nullable
+  private Map<String, Integer> _segmentToPartitionId;
+
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig, @Nullable ControllerMetrics controllerMetrics) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
     _tableConfig = tableConfig;
     _replication = tableConfig.getReplication();
     _partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
-    _controllerMetrics = controllerMetrics;
+    _partitionPinned = isPartitionPinned(tableConfig);
     if (_partitionColumn == null) {
-      _logger.info("Initialized with replication: {} without partition column for table: {} ", _replication,
-          _tableNameWithType);
+      _logger.info("Initialized with replication: {} without partition column, partitionPinned: {} for table: {}",
+          _replication, _partitionPinned, _tableNameWithType);
     } else {
-      _logger.info("Initialized with replication: {} and partition column: {} for table: {}", _replication,
-          _partitionColumn, _tableNameWithType);
+      _logger.info("Initialized with replication: {} and partition column: {}, partitionPinned: {} for table: {}",
+          _replication, _partitionColumn, _partitionPinned, _tableNameWithType);
+    }
+  }
+
+  /// Returns whether the segments of the table are pinned to their partition (see the class Javadoc): the table uses
+  /// strict replica-group routing. Subclasses add the conditions they need to pin the segments, and log a warning
+  /// when such a table cannot be pinned, in which case its segments are assigned with the segment assignment strategy.
+  protected boolean isPartitionPinned(TableConfig tableConfig) {
+    return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);

Review Comment:
   Deliberate, see the description: `multiStageReplicaGroup` will be removed. Added in 7df9c15 a warning at assignment initialization for an upsert/dedup table that is not pinned.



##########
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) {

Review Comment:
   Fixed in 7df9c15: the entries are stepped with the segments only when the CONSUMING segments of a real-time table are included. Otherwise they are set to the target at the first write, so a segment created mid-rebalance goes straight to the target instances, as before this PR.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1803,71 +1943,63 @@ private static void checkIfAnyServersAssignedMoreSegmentsThanBatchSize(int batch
     }
   }
 
-  /// Create a mapping of Pair(currentInstances, targetInstances) to partitionId to the current assignment of segments.
-  /// This is to be used for batching purposes for StrictReplicaGroup routing, for all segment assignment types:
-  /// RealtimeSegmentAssignment, StrictRealtimeSegmentAssignment and OfflineSegmentAssignment
-  /// @param currentAssignment the current assignment
-  /// @param targetAssignment the target assignment
-  /// @param segmentPartitionIdMap cache to store the partition ids to avoid fetching ZK segment metadata
-  /// @param partitionIdFetcher function to fetch the partition id
-  /// @return a mapping from Pair(currentInstances, targetInstances) to the partitionId to the segment assignment map of
-  ///         all segments that fall in that category
-  private static Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String, String>>>>
-      getCurrentAndTargetInstancesToPartitionIdToCurrentAssignmentMap(
-          Map<String, Map<String, String>> currentAssignment, Map<String, Map<String, String>> targetAssignment,
-          Object2IntOpenHashMap<String> segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher) {
-    Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String, String>>>>
-        currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap = new HashMap<>();
-
-    for (Map.Entry<String, Map<String, String>> assignment : currentAssignment.entrySet()) {
-      String segmentName = assignment.getKey();
-      Map<String, String> currentInstanceStateMap = assignment.getValue();
-      Map<String, String> targetInstanceStateMap = targetAssignment.get(segmentName);
-
-      int partitionId = segmentPartitionIdMap.computeIfAbsent(segmentName, v -> partitionIdFetcher.fetch(segmentName));
-      Pair<Set<String>, Set<String>> currentAndTargetInstances =
-          Pair.of(currentInstanceStateMap.keySet(), targetInstanceStateMap.keySet());
-      currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap.computeIfAbsent(currentAndTargetInstances,
-              k -> new TreeMap<>())
-          .computeIfAbsent(partitionId, k -> new TreeMap<>())
-          .put(segmentName, currentInstanceStateMap);
-    }
-
-    return currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap;
-  }
-
-  @VisibleForTesting
-  @FunctionalInterface
-  interface PartitionIdFetcher {
-    int fetch(String segmentName);
+  /// Segments and partitions with the same current and target instances, which are moved together for strict
+  /// replica-group routing.
+  private static class StrictReplicaGroup {
+    final Map<String, Map<String, String>> _segments = new TreeMap<>();
+    final List<Integer> _partitionIds = new ArrayList<>();
   }
 
-  private static class PartitionIdFetcherImpl implements PartitionIdFetcher {
-    private final String _tableNameWithType;
-    private final String _partitionColumn;
-    private final HelixManager _helixManager;
-    private final boolean _isStrictRealtimeSegmentAssignment;
-
-    private PartitionIdFetcherImpl(String tableNameWithType, @Nullable String partitionColumn,
-        HelixManager helixManager, boolean isStrictRealtimeSegmentAssignment) {
-      _tableNameWithType = tableNameWithType;
-      _partitionColumn = partitionColumn;
-      _helixManager = helixManager;
-      _isStrictRealtimeSegmentAssignment = isStrictRealtimeSegmentAssignment;
-    }
-
-    @Override
-    public int fetch(String segmentName) {
-      Integer partitionId =
-          SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType, _helixManager, _partitionColumn);
-      if (partitionId != null) {
-        return partitionId;
+  /// Groups the segments and the partitions by their current and target instances for batching with strict
+  /// replica-group routing, for all table types. Within a pair of current and target instances, the segments are
+  /// further grouped by the partition id in their name, so that partitions sharing the same instances can be moved in
+  /// different batches. When a segment has no partition id in its name (a segment of an offline table, or an uploaded
+  /// segment with a custom name), the whole pair is moved together so that the segment stays with the other segments
+  /// of its partition. The partitions of a partition-pinned table join the group of their current and target
+  /// instances. Partitions absent from the target partition assignment are dropped.
+  private static List<StrictReplicaGroup> getStrictReplicaGroups(Map<String, Map<String, String>> currentAssignment,
+      Map<String, Map<String, String>> targetAssignment, Map<Integer, List<String>> currentPartitionAssignment,
+      Map<Integer, List<String>> targetPartitionAssignment) {
+    Map<Pair<Set<String>, Set<String>>, Map<Integer, StrictReplicaGroup>> groupsMap = new HashMap<>();
+    for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) {
+      String segmentName = entry.getKey();
+      Map<String, String> currentInstanceStateMap = entry.getValue();
+      Pair<Set<String>, Set<String>> currentAndTargetInstances =
+          Pair.of(currentInstanceStateMap.keySet(), targetAssignment.get(segmentName).keySet());
+      Integer partitionId = SegmentUtils.getPartitionIdFromSegmentName(segmentName);

Review Comment:
   Fixed in 7df9c15: `rebalanceTable` returns the partition ids it resolves (one bulk ZK read), and the strict replica-group batching groups by the name-derived id, else the resolved id, else a hash of the name. The collapse into one group is gone. Covered in `TableRebalancerTest`.



-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332397690


##########
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:
   **Once a partition has an entry, the placement computed from the instance partitions is discarded, so evacuating servers via the instance partitions no longer redirects new segments.** Example: non-upsert strictReplicaGroup realtime table (newly pinned by this PR); partition p3's servers fail and mark the consuming replicas OFFLINE; the operator swaps them with `replaceInstance` without a full rebalance; ensureAllPartitionsConsuming -> createNewConsumingSegment -> assignSegment returns the entry, so the new CONSUMING segment goes back to the dead servers and p3 stays down until a rebalance (base: placed on the replacement servers). Uploads to qualifying offline tables behave the same. Consider letting the candidate win when the entry's instances are no longer in the instance partitions.



##########
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:
   **Every OFFLINE upload during a strictReplicaGroup rebalance forces a full recompute, which also misplaces tier-eligible uploads.** This check is gated on routing (`enableStrictReplicaGroup`), not on the table being pinned, and offline names never carry a partition id, so it returns null for each upload; the base took the copy path for offline tables. (1) The recompute reruns getInstancePartitionsMap (persisted with reassignInstances) and rebalanceTable; for a non-pinned replica-group table (numPartitions>1) the strategy does one uncached ZK read per segment, and with continuous pushes the versioned write keeps losing (ZkBadVersion only rolls back and loops), so the rebalance can stall. (2) The recompute reuses the tier cache frozen at rebalance start (`_segmentToTierName`; TierSegmentAssignment treats absent segments as non-tier), so a tier-eligible upload that tiered segment assignment placed on tier servers is retargeted to the default instances, then moved back by the next Seg
 mentRelocator run.



##########
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:
   **A segment added mid-rebalance to a partition without an entry is left in place, whereas the base always fully recomputed for strict (upsert/dedup) tables.** Example (rolling upgrade): pre-PR upsert table with no entries; a new controller rebalances p0 [A,B]->[C,D]; the old-version table lead resumes/repairs p0's consumption or accepts an upload for p0 and places the new segment on [A,B] without creating an entry; no segment in segmentsToMove changes state, so this branch keeps it at [A,B], the rest of p0 moves to [C,D], the final step sets entry p0=[C,D], and the rebalance reports DONE with the upsert partition split.



##########
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:
   **When a partition has no entry and no ONLINE/CONSUMING LLC segment, this reads the ZK metadata of every custom-named segment one by one inside the IdealState updater, now for dedup and plain strictReplicaGroup tables too** (base: only upsert scanned; dedup skipped uploaded segments; plain tables never scanned). Example: plain/dedup strictReplicaGroup table with N MergeRollup `merged_*` segments gains K stream partitions; the next ensureAllPartitionsConsuming updater does K x N sequential ZK reads per attempt (up to 10 attempts) while holding the table's group-commit queue, delaying segment-commit updates; an orphan segment in the IdealState without ZK metadata makes `SegmentUtils.getSegmentPartitionId` throw and the new partitions never start consuming. Checking relocation first and batch-reading only the needed segments outside the updater would avoid 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());

Review Comment:
   **putIfAbsent cannot tell a partition added while waiting from one the target deliberately drops, so leftover entries can survive the rebalance.** Example: formerly pinned OFFLINE strictReplicaGroup table whose instance partitions now hold 2 instances per partition (rebalanceTable returns `Map.of()`), entries {0:[A],1:[B]}; a retention deletion of a non-moving segment lands before the first write, this makes target == current, and the rebalance ends DONE with the stale entries; if the instance partitions are later set back to one instance per partition via assignInstances/PUT without a rebalance, uploads for partition 0 go to stale server A. Only re-adding partitions that are new since the target was computed would avoid this.



-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332393306


##########
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",

Review Comment:
   **Realtime strictReplicaGroup tables that are not upsert/dedup now hard-fail on segments without a derivable partition id.** They are partition-pinned by this PR, and both `assignPinnedConsumingSegment` (here) and `getTargetInstances` (line 407) throw `Failed to find partition id`, where the base fell back to `getSegmentPartitionIdOrDefault` (hash). Example: uploading a custom-named segment or a MergeRollupTask `merged_...` output (supported for non-upsert realtime tables) to such a table without segmentPartitionConfig fails inside the IdealState updater and its ZK metadata is deleted; if such a segment already exists, every rebalance (REST, tenant, SegmentRelocator) returns FAILED. `OfflineSegmentAssignment` falls back to the strategy in the same case.



-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332399450


##########
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:
   **A new ONLINE segment sitting on its partition entry is retargeted to the partition's CONSUMING-placement target even when the table relocates COMPLETED segments**, which disagrees with what `rebalanceTable` computes (COMPLETED strategy). Example: non-upsert strictReplicaGroup table with COMPLETED instance partitions, rebalanced with includeConsuming=false; during one EV wait partition P commits twice, so segment S created and committed within that wait is ONLINE on P's intermediate entry and absent from the target; this line sends it to P's consuming target, leaving a completed segment on consuming servers (extra moves, fixed only by a later relocation).



##########
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:
   **Data-less partition entries take part in the min-available-replicas accounting like real segments.** `canAssignWithStrictReplicaGroup` registers an entry's available instances in the shared `availableInstancesMap`, so with batching on (groups iterated in HashMap order) an entry-only group processed first can hold back real segments heading to the same next instance set. Example: minAvailableReplicas=1, partition 1's entry [p,q]->[p,r] with no segments in that pair, partition 3's segments [r,s]->[p,r]; both step to [p,r] with available {p} vs {r}; if partition 1's group comes first, partition 3's segments fail `retainAll` and stay put for an extra round although no real segment constrains them. Deriving an entry's next step from its segments' next assignment would avoid this.



##########
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:
   **OfflineSegmentAssignment does not override `getExistingAssignment`, so for offline tables created before this PR a partition's first entry is seeded from the current instance partitions rather than from where its segments are.** That contradicts the class doc ('even when the instance partitions have changed'), and the realtime side has this handling. Example: pre-existing pinned OFFLINE (e.g. upsert) table without entries: replaceInstance X->Y without a rebalance, then the first upload of partition P stores entry [Y]; the operator reverts the instance partitions to X, but every later upload of P still goes to Y (base: back to X), so P stays split with a growing share on Y until a rebalance.



##########
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:
   **Any list field whose key passes lenient `Integer.parseInt` is treated as a partition entry, including negative keys.** A hand-edited `-1` key makes every `rebalanceTable` throw while seeding the target: `OfflineSegmentAssignment.getPartitionInstances(-1)` computes `-1 % numPartitions = -1`, so `getInstances(-1, 0)` is null and NPEs; `RealtimeSegmentAssignment.assignConsumingSegment(-1, ...)` throws IndexOutOfBounds. The rebalance returns FAILED and nothing removes the key. A JSON-null list value similarly NPEs at `new HashSet<>(entry.getValue())`. Only canonical non-negative keys (`Integer.toString(id)`, as `getInstances` writes) should be accepted.



-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332395498


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   @Nullable
   private Map<String, String> _segmentToTierName;
 
+  // Cache of segment name to partition id (null value means the partition id cannot be determined) for the segments
+  // whose partition id is not in the segment name, loaded with a single bulk ZK read on first use. Same scope and
+  // limitations as _segmentToTierName; a segment absent from the cache is resolved on demand and merged in.
+  @Nullable
+  private Map<String, Integer> _segmentToPartitionId;
+
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig, @Nullable ControllerMetrics controllerMetrics) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
     _tableConfig = tableConfig;
     _replication = tableConfig.getReplication();
     _partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
-    _controllerMetrics = controllerMetrics;
+    _partitionPinned = isPartitionPinned(tableConfig);
     if (_partitionColumn == null) {
-      _logger.info("Initialized with replication: {} without partition column for table: {} ", _replication,
-          _tableNameWithType);
+      _logger.info("Initialized with replication: {} without partition column, partitionPinned: {} for table: {}",
+          _replication, _partitionPinned, _tableNameWithType);
     } else {
-      _logger.info("Initialized with replication: {} and partition column: {} for table: {}", _replication,
-          _partitionColumn, _tableNameWithType);
+      _logger.info("Initialized with replication: {} and partition column: {}, partitionPinned: {} for table: {}",
+          _replication, _partitionColumn, _partitionPinned, _tableNameWithType);
+    }
+  }
+
+  /// Returns whether the segments of the table are pinned to their partition (see the class Javadoc): the table uses
+  /// strict replica-group routing. Subclasses add the conditions they need to pin the segments, and log a warning
+  /// when such a table cannot be pinned, in which case its segments are assigned with the segment assignment strategy.
+  protected boolean isPartitionPinned(TableConfig tableConfig) {
+    return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);
+  }
+
+  /// Returns the instances of the given partition of a partition-pinned table, creating its entry in the partition
+  /// assignment of the given ideal state when the partition has no entry yet, or an empty one (a rebalance in low disk
+  /// mode without minimum available replicas empties the entry while it drops the partition from its current
+  /// instances). The entry is created from the existing assignment of the partition (see [#getExistingAssignment])
+  /// when there is one, else from the given candidate instances computed from the instance partitions. An existing
+  /// entry wins over the candidate instances, so that the segments of the partition stay collocated when the instance
+  /// partitions change: the partition moves on the next rebalance.
+  protected List<String> getOrCreatePartitionInstances(IdealState idealState, int partitionId,
+      List<String> candidateInstances) {
+    List<String> instances = PartitionAssignmentUtils.getInstances(idealState, partitionId);
+    if (CollectionUtils.isEmpty(instances)) {

Review Comment:
   **A non-empty entry is trusted unconditionally, but only the new TableRebalancer keeps entries in sync with segment moves.** Entries are written on the first commit per partition (not only on rebalance), and any move made without the new rebalancer leaves them stale; they then redirect every new segment of the partition. Example: a new-version controller writes entry `0`=[S1,S2]; controllers are rolled back (or an old-version controller runs a rebalance during a rolling upgrade) and the old TableRebalancer, which only calls `setMapFields`, moves partition 0 to [S2,S3]; after (re-)upgrade every new CONSUMING segment of partition 0 goes to [S1,S2] (S1 possibly decommissioned) while its history is on [S2,S3], splitting the upsert/dedup partition until a new-version rebalance. The removed mismatch meter was the only signal. CLAUDE.md asks to preserve mixed-version compatibility; consider validating the entry against the partition's existing segments (e.g. its CONSUMING segment) before
  trusting it.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   @Nullable
   private Map<String, String> _segmentToTierName;
 
+  // Cache of segment name to partition id (null value means the partition id cannot be determined) for the segments
+  // whose partition id is not in the segment name, loaded with a single bulk ZK read on first use. Same scope and
+  // limitations as _segmentToTierName; a segment absent from the cache is resolved on demand and merged in.
+  @Nullable
+  private Map<String, Integer> _segmentToPartitionId;
+
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig, @Nullable ControllerMetrics controllerMetrics) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
     _tableConfig = tableConfig;
     _replication = tableConfig.getReplication();
     _partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
-    _controllerMetrics = controllerMetrics;
+    _partitionPinned = isPartitionPinned(tableConfig);
     if (_partitionColumn == null) {
-      _logger.info("Initialized with replication: {} without partition column for table: {} ", _replication,
-          _tableNameWithType);
+      _logger.info("Initialized with replication: {} without partition column, partitionPinned: {} for table: {}",
+          _replication, _partitionPinned, _tableNameWithType);
     } else {
-      _logger.info("Initialized with replication: {} and partition column: {} for table: {}", _replication,
-          _partitionColumn, _tableNameWithType);
+      _logger.info("Initialized with replication: {} and partition column: {}, partitionPinned: {} for table: {}",
+          _replication, _partitionColumn, _partitionPinned, _tableNameWithType);
+    }
+  }
+
+  /// Returns whether the segments of the table are pinned to their partition (see the class Javadoc): the table uses
+  /// strict replica-group routing. Subclasses add the conditions they need to pin the segments, and log a warning
+  /// when such a table cannot be pinned, in which case its segments are assigned with the segment assignment strategy.
+  protected boolean isPartitionPinned(TableConfig tableConfig) {
+    return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);

Review Comment:
   **Upsert/dedup tables on `multiStageReplicaGroup` routing silently lose partition pinning.** `TableConfigUtils.UPSERT_DEDUP_ALLOWED_ROUTING_STRATEGIES` still accepts that routing, and the deleted factory pinned every upsert/dedup table regardless of routing. Example: such a table's CONSUMING instance partitions change without moving segments (assignInstances, PUT instancePartitions, replaceInstance, or a reassignInstances rebalance that persists new instance partitions and then fails/stops); the next commit of partition P runs `assignConsumingSegment` from the instance partitions and puts the new CONSUMING segment on the new server while P's completed segments stay on the old one, so upsert serves stale and new rows for the same key and dedup admits duplicates. The PR calls this deliberate, but nothing rejects, warns about or migrates the still-valid config.



##########
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) {

Review Comment:
   **With includeConsuming=false on a pinned non-upsert table, CONSUMING segments created mid-rebalance are stranded on intermediate instances.** The partition entry still moves step by step (its target comes from the instance partitions), new CONSUMING segments are assigned to the entry at that moment, and both this `continue` and the full recompute (`newAssignment.putAll(consumingSegmentAssignment)`) leave them there. Example: non-upsert strictReplicaGroup table rebalanced [A,B]->[C,D] with includeConsuming=false (RebalanceConfig and SegmentRelocator default) and minAvailableReplicas=1: step 1 writes entry [A,C] ([A] in lowDiskMode); p0__6 is created CONSUMING on {A,C}; the rebalance reports DONE with p0__6 on A (being decommissioned) or with a single replica, and after it commits it stays split from the rest of p0 on [C,D]. Before the PR the new segment went straight to the new instance partitions {C,D}.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1803,71 +1943,63 @@ private static void checkIfAnyServersAssignedMoreSegmentsThanBatchSize(int batch
     }
   }
 
-  /// Create a mapping of Pair(currentInstances, targetInstances) to partitionId to the current assignment of segments.
-  /// This is to be used for batching purposes for StrictReplicaGroup routing, for all segment assignment types:
-  /// RealtimeSegmentAssignment, StrictRealtimeSegmentAssignment and OfflineSegmentAssignment
-  /// @param currentAssignment the current assignment
-  /// @param targetAssignment the target assignment
-  /// @param segmentPartitionIdMap cache to store the partition ids to avoid fetching ZK segment metadata
-  /// @param partitionIdFetcher function to fetch the partition id
-  /// @return a mapping from Pair(currentInstances, targetInstances) to the partitionId to the segment assignment map of
-  ///         all segments that fall in that category
-  private static Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String, String>>>>
-      getCurrentAndTargetInstancesToPartitionIdToCurrentAssignmentMap(
-          Map<String, Map<String, String>> currentAssignment, Map<String, Map<String, String>> targetAssignment,
-          Object2IntOpenHashMap<String> segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher) {
-    Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String, String>>>>
-        currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap = new HashMap<>();
-
-    for (Map.Entry<String, Map<String, String>> assignment : currentAssignment.entrySet()) {
-      String segmentName = assignment.getKey();
-      Map<String, String> currentInstanceStateMap = assignment.getValue();
-      Map<String, String> targetInstanceStateMap = targetAssignment.get(segmentName);
-
-      int partitionId = segmentPartitionIdMap.computeIfAbsent(segmentName, v -> partitionIdFetcher.fetch(segmentName));
-      Pair<Set<String>, Set<String>> currentAndTargetInstances =
-          Pair.of(currentInstanceStateMap.keySet(), targetInstanceStateMap.keySet());
-      currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap.computeIfAbsent(currentAndTargetInstances,
-              k -> new TreeMap<>())
-          .computeIfAbsent(partitionId, k -> new TreeMap<>())
-          .put(segmentName, currentInstanceStateMap);
-    }
-
-    return currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap;
-  }
-
-  @VisibleForTesting
-  @FunctionalInterface
-  interface PartitionIdFetcher {
-    int fetch(String segmentName);
+  /// Segments and partitions with the same current and target instances, which are moved together for strict
+  /// replica-group routing.
+  private static class StrictReplicaGroup {
+    final Map<String, Map<String, String>> _segments = new TreeMap<>();
+    final List<Integer> _partitionIds = new ArrayList<>();
   }
 
-  private static class PartitionIdFetcherImpl implements PartitionIdFetcher {
-    private final String _tableNameWithType;
-    private final String _partitionColumn;
-    private final HelixManager _helixManager;
-    private final boolean _isStrictRealtimeSegmentAssignment;
-
-    private PartitionIdFetcherImpl(String tableNameWithType, @Nullable String partitionColumn,
-        HelixManager helixManager, boolean isStrictRealtimeSegmentAssignment) {
-      _tableNameWithType = tableNameWithType;
-      _partitionColumn = partitionColumn;
-      _helixManager = helixManager;
-      _isStrictRealtimeSegmentAssignment = isStrictRealtimeSegmentAssignment;
-    }
-
-    @Override
-    public int fetch(String segmentName) {
-      Integer partitionId =
-          SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType, _helixManager, _partitionColumn);
-      if (partitionId != null) {
-        return partitionId;
+  /// Groups the segments and the partitions by their current and target instances for batching with strict
+  /// replica-group routing, for all table types. Within a pair of current and target instances, the segments are
+  /// further grouped by the partition id in their name, so that partitions sharing the same instances can be moved in
+  /// different batches. When a segment has no partition id in its name (a segment of an offline table, or an uploaded
+  /// segment with a custom name), the whole pair is moved together so that the segment stays with the other segments
+  /// of its partition. The partitions of a partition-pinned table join the group of their current and target
+  /// instances. Partitions absent from the target partition assignment are dropped.
+  private static List<StrictReplicaGroup> getStrictReplicaGroups(Map<String, Map<String, String>> currentAssignment,
+      Map<String, Map<String, String>> targetAssignment, Map<Integer, List<String>> currentPartitionAssignment,
+      Map<Integer, List<String>> targetPartitionAssignment) {
+    Map<Pair<Set<String>, Set<String>>, Map<Integer, StrictReplicaGroup>> groupsMap = new HashMap<>();
+    for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) {
+      String segmentName = entry.getKey();
+      Map<String, String> currentInstanceStateMap = entry.getValue();
+      Pair<Set<String>, Set<String>> currentAndTargetInstances =
+          Pair.of(currentInstanceStateMap.keySet(), targetAssignment.get(segmentName).keySet());
+      Integer partitionId = SegmentUtils.getPartitionIdFromSegmentName(segmentName);

Review Comment:
   **batchSizePerServer no longer limits OFFLINE strictReplicaGroup rebalances.** The partition id comes from the segment name only; offline names never carry one, so every segment maps to UNKNOWN_PARTITION_ID and the merge below moves the whole (current,target) pair as one group, which is always admitted because the first group reaching a server has `segmentsAddedToServerSoFar == 0`. Example: offline strictReplicaGroup table with replica-group instance partitions, batchSizePerServer=100, replacing server X with X': all ~2,500 segments of the pair ({X,Y},{X',Y}) are added to X' in one IdealState step (only a WARN is logged), where the removed PartitionIdFetcher grouped by ZK partition id or a per-segment hash and moved ~100 per step. One custom-named segment does the same to a realtime pinned pair. The partition ids that rebalanceTable already resolved (`_segmentToPartitionId`) could be returned and used here, with a per-segment default for unpinned tables.



##########
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:
   **The partition-only update writes with the version read at the start of doRebalance and turns a conflict into FAILED with no retry; the dry-run meanwhile reports NO_OP.** Between the read (~287) and this write run instance-partition computation, rebalanceTable (bulk ZK reads), fetchTableSizeDetails (up to 30s) and the summary, so on a busy table any commit makes `set(...)` throw ZkBadVersionException and the catch returns FAILED, where the base returned NO_OP. Example: pre-PR upsert table, segments balanced, an idle or paused partition without entry: a tenant rebalance or REST downtime=true rebalance fails and every rerun races again. The default REST flow runs a dry-run first, gets NO_OP and stops, so the PR's 'first rebalance populates them' never happens through it. Consider routing partition-only changes through the retrying main loop.



-- 
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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
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: 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


Re: [PR] Track partition assignment in IdealState for strict replica-group tables [pinot]

Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on PR #19916:
URL: https://github.com/apache/pinot/pull/19916#issuecomment-6388866439

   Documentation follow-up merged: https://github.com/pinot-contrib/pinot-docs/pull/1143 — explains persistent partition placement, rebalance migration, in-use instance references and removal of the assignment mismatch meter.


-- 
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