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


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, 
Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = 
_assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre 
conditions state would return false
+    if (segmentAssignmentStrategy instanceof 
DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = 
instancePartitionsMap.get(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 = assignSegment(segmentName, 
currentAssignment,
+            instancePartitions, segmentAssignmentStrategy, 
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);
+    SegmentAssignmentStrategy segmentAssignmentStrategy = 
_assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre 
conditions state would return false
+    if (segmentAssignmentStrategy instanceof 
DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.reassignSegments(currentAssignment,
+          offlineInstancePartitions, InstancePartitionsType.OFFLINE);

Review Comment:
   For Dim tables, we would have reassignSegments bypassing bootstrap, so 
should be fine



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