somandal commented on code in PR #15817:
URL: https://github.com/apache/pinot/pull/15817#discussion_r2116356028


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -439,21 +441,66 @@ public Map<String, Map<String, String>> 
getNonTierSegmentAssignment() {
     }
   }
 
+  /**
+   * Return the partitionId for an OFFLINE or COMPLETED instance partitions of 
a REALTIME table with relocation enabled
+   * The partitionId will be calculated as:
+   * <ul>
+   *   <li>
+   *     1. If numPartitions = 1, return partitionId = 0
+   *   </li>
+   *   <li>
+   *     2. Otherwise, fallback to either the OFFLINE or REALTIME partitionId 
calculation logic
+   *   </li>
+   * </ul>
+   */
+  public static int getOfflineOrCompletedPartitionId(String segmentName, 
String tableName, TableType tableType,
+      HelixManager helixManager, int numPartitions, @Nullable String 
partitionColumn) {
+    int partitionId;
+    if (numPartitions == 1) {
+      partitionId = 0;
+    } else {
+      // Uniformly spray the segment partitions over the instance partitions
+      if (tableType == TableType.OFFLINE) {
+        partitionId = SegmentAssignmentUtils
+                .getOfflineSegmentPartitionId(segmentName, tableName, 
helixManager, partitionColumn);
+      } else {
+        partitionId = SegmentAssignmentUtils
+            .getRealtimeSegmentPartitionId(segmentName, tableName, 
helixManager, partitionColumn);
+      }
+    }
+    return partitionId;
+  }
+
   /**
    * Returns a partition id for offline table
+   * The partitionId will be calculated as:
+   * <ul>
+   *   <li>
+   *     1. If partitionColumn == null, return a default partitionId 
calculated based on the hashCode of the segmentName
+   *   </li>
+   *   <li>
+   *     2. Otherwise, fetch the partitionMetadata from the SegmentZkMetadata 
related to the partitionColumn and return
+   *        that as the partitionId
+   *   </li>
+   * </ul>
    */
   public static int getOfflineSegmentPartitionId(String segmentName, String 
offlineTableName, HelixManager helixManager,
       @Nullable String partitionColumn) {
     SegmentZKMetadata segmentZKMetadata =
         
ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), 
offlineTableName, segmentName);
     Preconditions.checkState(segmentZKMetadata != null,
         "Failed to find segment ZK metadata for segment: %s of table: %s", 
segmentName, offlineTableName);
-    return getPartitionId(segmentZKMetadata, offlineTableName, 
partitionColumn);
+    return getOfflinePartitionId(segmentZKMetadata, offlineTableName, 
partitionColumn);
   }
 
-  private static int getPartitionId(SegmentZKMetadata segmentZKMetadata, 
String offlineTableName,
+  private static int getOfflinePartitionId(SegmentZKMetadata 
segmentZKMetadata, String offlineTableName,

Review Comment:
   done, moved this one out. let me know if you want me to move the others too



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

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


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

Reply via email to