GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r971256430


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -53,62 +49,57 @@ protected int getReplication(TableConfig tableConfig) {
   public List<String> assignSegment(String segmentName, Map<String, 
Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     InstancePartitions instancePartitions = 
instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = 
SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, 
instancePartitions, InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find 
OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for 
table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, 
currentAssignment, instancePartitions);
+    List<String> instancesAssigned = 
segmentAssignmentStrategy.assignSegment(segmentName,
+        currentAssignment, instancePartitions, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", 
segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        
ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), 
_tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", 
segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        
segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain 
partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple 
partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, 
Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, 
@Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, 
Configuration config) {
     InstancePartitions offlineInstancePartitions = 
instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = 
SegmentAssignmentStrategyFactory

Review Comment:
   we already have that, added the same preconditions check before



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to