[GitHub] [pinot] klsince commented on a diff in pull request #10746: [feature] Consider tierConfigs when assigning new offline segment

["klsince (via GitHub)" <gi...@apache.org>]

klsince commented on code in PR #10746:
URL: https://github.com/apache/pinot/pull/10746#discussion_r1187869591


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/TierConfigUtils.java:
##########
@@ -63,6 +68,33 @@ public static String getDataDirForTier(TableConfig tableConfig, String tierName)
     return getDataDirForTier(tableConfig, tierName, Collections.emptyMap());
   }
 
+  /**
+   * Compute default instance partitions for every configured tier
+   *
+   * @return a map with tier names as keys, and default instance partitions as values
+   */
+  public static Map<String, InstancePartitions> getTierToInstancePartitionsMap(String tableNameWithType,

Review Comment:
   try to reuse this for TableRebalancer.getTierToInstancePartitionsMap?



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/TierConfigUtils.java:
##########
@@ -63,6 +68,33 @@ public static String getDataDirForTier(TableConfig tableConfig, String tierName)
     return getDataDirForTier(tableConfig, tierName, Collections.emptyMap());
   }
 
+  /**
+   * Compute default instance partitions for every configured tier
+   *
+   * @return a map with tier names as keys, and default instance partitions as values
+   */
+  public static Map<String, InstancePartitions> getTierToInstancePartitionsMap(String tableNameWithType,
+      @Nullable List<Tier> sortedTiers, HelixManager helixManager) {
+    if (sortedTiers == null) {
+      return null;
+    }
+
+    Map<String, InstancePartitions> tierToInstancePartitionsMap = new HashMap<>();
+    for (Tier tier : sortedTiers) {
+      LOGGER.info("Fetching/computing instance partitions for tier: {} of table: {}", tier.getName(),
+          tableNameWithType);
+
+      final PinotServerTierStorage storage = (PinotServerTierStorage) tier.getStorage();

Review Comment:
   no need to use `final` according to current coding convention in this repo, although it's good habit.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -408,6 +409,25 @@ public Map<String, Map<String, String>> getNonTierSegmentAssignment() {
     }
   }
 
+  /**
+   * Checks tiers one-by-one to return the first one for which a newly registered segment qualifies
+   * @param tableNameWithType table to which the segment assignment belongs
+   * @param sortedTiers list of tiers, pre-sorted as per desired order by caller
+   * @param segmentName name of the new segment
+   * @return tier for which the segment qualifies, or null if segment doesn't belong to any tier
+   */
+  public static Tier findTierForNewSegment(String tableNameWithType, @Nullable List<Tier> sortedTiers,

Review Comment:
   move this to TierConfigUtils.java instead?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java:
##########
@@ -2268,10 +2271,28 @@ public void assignTableSegment(String tableNameWithType, String segmentName) {
     try {
       TableConfig tableConfig = getTableConfig(tableNameWithType);
       Preconditions.checkState(tableConfig != null, "Failed to find table config for table: " + tableNameWithType);
+
+      List<Tier> sortedTiers = null;
+      Map<String, InstancePartitions> tierToInstancePartitionsMap = null;
+
+      // Initialize tier information only in case direct tier assignment is configured
+      if (_enableTieredSegmentAssignment && CollectionUtils.isNotEmpty(tableConfig.getTierConfigsList())) {
+        sortedTiers = TierConfigUtils.getSortedTiersForStorageType(tableConfig.getTierConfigsList(),
+            TierFactory.PINOT_SERVER_STORAGE_TYPE, _helixZkManager);
+        tierToInstancePartitionsMap =
+            TierConfigUtils.getTierToInstancePartitionsMap(tableNameWithType, sortedTiers, _helixZkManager);
+
+        // Update segment tier to support direct assignment for multiple data directories
+        updateSegmentTargetTier(tableNameWithType, segmentName, sortedTiers);

Review Comment:
   we can move this stmt ahead of this if-block to get a default `instancePartitionsMap`, 
   ```
   Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap =
             fetchOrComputeInstancePartitions(tableNameWithType, tableConfig);
   ```
   
   then, in this if-block, we can figure out the `tierInstancePartitions` and overwrite `instancePartitionsMap` if there is a valid `tierInstancePartitions`
   ```
   if (_enableTieredSegmentAssignment && CollectionUtils.isNotEmpty(tableConfig.getTierConfigsList())) {
   ...
     if (tierInstancePartitions != null) {
       instancePartitionsMap = tierInstancePartitions
     }
   ...
   }
   ```



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -32,19 +32,63 @@
 import org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategyFactory;
 import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
 import org.apache.pinot.spi.utils.RebalanceConfigConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
  * Segment assignment for offline table.
  */
 public class OfflineSegmentAssignment extends BaseSegmentAssignment {
+  private final Logger _logger = LoggerFactory.getLogger(getClass());
 
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    // Fallback to default assignment
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
+    return doAssignSegment(segmentName, currentAssignment, instancePartitions);
+  }
+
+  @Override
+  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
+      @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap) {

Review Comment:
   It looks like we don't have to overload this assignSegment(). 
   
   We can just add a helper method in PinotHelixResourceManager class to figure out the `tierInstancePartitions` and call the original assignSegment() method, saving all the changes on SegmentAssignment, OfflineSegmentAssignment and RealtimeSegmentAssignment. 



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