Cleanup segment assignment strategy interface (#17038)

Author: yashmayya

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java

@@ -33,7 +33,6 @@
 import org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategyFactory;
 import org.apache.pinot.segment.local.utils.TableConfigUtils;
 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;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -94,8 +93,7 @@ public void init(HelixManager helixManager, TableConfig tableConfig, @Nullable C
    */
   protected Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, String>>> rebalanceTiers(
       Map<String, Map<String, String>> currentAssignment, @Nullable List<Tier> sortedTiers,
-      @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, boolean bootstrap,
-      InstancePartitionsType instancePartitionsType) {
+      @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, boolean bootstrap) {
     if (sortedTiers == null) {
       return Pair.of(null, currentAssignment);
     }
@@ -129,7 +127,7 @@ protected Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, S
       _logger.info("Rebalancing tier: {} for table: {} with bootstrap: {}, instance partitions: {}", tierName,
           _tableNameWithType, bootstrap, tierInstancePartitions);
       newTierAssignments.add(reassignSegments(tierName, tierCurrentAssignment, tierInstancePartitions, bootstrap,
-          segmentAssignmentStrategy, instancePartitionsType));
+          segmentAssignmentStrategy));
     }
 
     return Pair.of(newTierAssignments, tierSegmentAssignment.getNonTierSegmentAssignment());
@@ -140,7 +138,7 @@ protected Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, S
    */
   protected Map<String, Map<String, String>> reassignSegments(String instancePartitionType,
       Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, boolean bootstrap,
-      SegmentAssignmentStrategy segmentAssignmentStrategy, InstancePartitionsType instancePartitionsType) {
+      SegmentAssignmentStrategy segmentAssignmentStrategy) {
     Map<String, Map<String, String>> newAssignment;
     if (bootstrap) {
       _logger.info("Bootstrapping segment assignment for {} segments of table: {}", instancePartitionType,
@@ -150,14 +148,14 @@ protected Map<String, Map<String, String>> reassignSegments(String instanceParti
       newAssignment = new TreeMap<>();
       for (String segment : currentAssignment.keySet()) {
         List<String> assignedInstances =
-            segmentAssignmentStrategy.assignSegment(segment, newAssignment, instancePartitions, instancePartitionsType);
+            segmentAssignmentStrategy.assignSegment(segment, newAssignment, instancePartitions);
         newAssignment
             .put(segment, SegmentAssignmentUtils.getInstanceStateMap(assignedInstances, SegmentStateModel.ONLINE));
       }
     } else {
       // Use segment assignment strategy
       newAssignment =
-          segmentAssignmentStrategy.reassignSegments(currentAssignment, instancePartitions, instancePartitionsType);
+          segmentAssignmentStrategy.reassignSegments(currentAssignment, instancePartitions);
     }
     return newAssignment;
   }

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java

@@ -51,8 +51,7 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
     List<String> instancesAssigned =
-        segmentAssignmentStrategy.assignSegment(segmentName, currentAssignment, instancePartitions,
-            InstancePartitionsType.OFFLINE);
+        segmentAssignmentStrategy.assignSegment(segmentName, currentAssignment, instancePartitions);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
@@ -74,22 +73,20 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     // tierPartitionMap has single tier for Dim tables and remove below check
     // See https://github.com/apache/pinot/issues/9047
     if (segmentAssignmentStrategy instanceof AllServersSegmentAssignmentStrategy) {
-      return segmentAssignmentStrategy.reassignSegments(currentAssignment, offlineInstancePartitions,
-          InstancePartitionsType.OFFLINE);
+      return segmentAssignmentStrategy.reassignSegments(currentAssignment, offlineInstancePartitions);
     }
     boolean bootstrap = config.isBootstrap();
     // Rebalance tiers first
     Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, String>>> pair =
-        rebalanceTiers(currentAssignment, sortedTiers, tierInstancePartitionsMap, bootstrap,
-            InstancePartitionsType.OFFLINE);
+        rebalanceTiers(currentAssignment, sortedTiers, tierInstancePartitionsMap, bootstrap);
     List<Map<String, Map<String, String>>> newTierAssignments = pair.getLeft();
     Map<String, Map<String, String>> nonTierAssignment = pair.getRight();
 
     _logger.info("Rebalancing table: {} with instance partitions: {}, bootstrap: {}", _tableNameWithType,
         offlineInstancePartitions, bootstrap);
     Map<String, Map<String, String>> newAssignment =
         reassignSegments(InstancePartitionsType.OFFLINE.toString(), nonTierAssignment, offlineInstancePartitions,
-            bootstrap, segmentAssignmentStrategy, InstancePartitionsType.OFFLINE);
+            bootstrap, segmentAssignmentStrategy);
 
     // Add tier assignments, if available
     if (CollectionUtils.isNotEmpty(newTierAssignments)) {

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java

@@ -96,8 +96,7 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
       SegmentAssignmentStrategy segmentAssignmentStrategy =
           SegmentAssignmentStrategyFactory.getSegmentAssignmentStrategy(_helixManager, _tableConfig,
               instancePartitionsType.toString(), instancePartitions);
-      instancesAssigned = segmentAssignmentStrategy.assignSegment(segmentName, currentAssignment, instancePartitions,
-          InstancePartitionsType.COMPLETED);
+      instancesAssigned = segmentAssignmentStrategy.assignSegment(segmentName, currentAssignment, instancePartitions);
     } else {
       instancesAssigned = assignConsumingSegment(segmentName, instancePartitions);
     }
@@ -186,8 +185,7 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap = config.isBootstrap();
     // Rebalance tiers first
     Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, String>>> pair =
-        rebalanceTiers(currentAssignment, sortedTiers, tierInstancePartitionsMap, bootstrap,
-            InstancePartitionsType.COMPLETED);
+        rebalanceTiers(currentAssignment, sortedTiers, tierInstancePartitionsMap, bootstrap);
 
     List<Map<String, Map<String, String>>> newTierAssignments = pair.getLeft();
     Map<String, Map<String, String>> nonTierAssignment = pair.getRight();
@@ -220,7 +218,7 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
       _logger.info("Reassigning COMPLETED segments with COMPLETED instance partitions for table: {}",
           _tableNameWithType);
       newAssignment = reassignSegments(InstancePartitionsType.COMPLETED.toString(), completedSegmentAssignment,
-          completedInstancePartitions, bootstrap, segmentAssignmentStrategy, InstancePartitionsType.COMPLETED);
+          completedInstancePartitions, bootstrap, segmentAssignmentStrategy);
     } else {
       // When COMPLETED instance partitions are not provided, reassign COMPLETED segments the same way as CONSUMING
       // segments with CONSUMING instance partitions (ensure COMPLETED segments are served by the correct instances when

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/AllServersSegmentAssignmentStrategy.java

@@ -31,7 +31,6 @@
 import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TenantConfig;
-import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -41,34 +40,30 @@
  * Segment assignment for an offline dimension table.
  * <ul>
  *   <li>
- *     <p>This segment assignment strategy is used when {@link TableConfig#IS_DIM_TABLE_KEY}is
- *     set to "true".</p>
+ *     This segment assignment strategy is used when {@link TableConfig#IS_DIM_TABLE_KEY} is set to "true".
  *   </li>
  *   <li>
- *     <p>For a dimension table we assign the segment to all the hosts. Thus for this assignment
- *     strategy we simply return all the hosts under a given tag as the assigned hosts for
- *     a given segment.</p>
+ *     For a dimension table we assign all segments to all the servers. Thus, for this assignment strategy, we simply
+ *     return all the instances under a given tag as the assigned instances for any given segment.
  *   </li>
  * </ul>
  */
 public class AllServersSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
   private static final Logger LOGGER = LoggerFactory.getLogger(AllServersSegmentAssignmentStrategy.class);
 
   private HelixManager _helixManager;
-  private String _offlineTableName;
   private TenantConfig _tenantConfig;
 
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig) {
     _helixManager = helixManager;
-    _offlineTableName = tableConfig.getTableName();
     _tenantConfig = tableConfig.getTenantConfig();
-    LOGGER.info("Initialized AllServersSegmentAssignmentStrategy for table: {}", _offlineTableName);
+    LOGGER.info("Initialized AllServersSegmentAssignmentStrategy for table: {}", tableConfig.getTableName());
   }
 
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
     String serverTag = _tenantConfig.getServer();
     Set<String> instances = HelixHelper.getServerInstancesForTenant(_helixManager, serverTag);
     int numInstances = instances.size();
@@ -79,7 +74,7 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
 
   @Override
   public Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
 
     String serverTag = _tenantConfig.getServer();
     Set<String> instances = HelixHelper.getServerInstancesForTenant(_helixManager, serverTag);

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalancedNumSegmentAssignmentStrategy.java

@@ -26,18 +26,15 @@
 import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
 import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
-import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 
 /**
- * Balance num Segment assignment strategy class for offline segment assignment
- * <ul>
- *   <li>
- *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
- *   </li>
- * </ul>
+ * Balanced segment assignment strategy class where segments are distributed to instances such that each instance has
+ * approximately the same number of segments. This is the default segment assignment strategy.
+ * <p>
+ * This segment assignment strategy is used when table replication/ num_replica_groups = 1.
  */
 public class BalancedNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
   private static final Logger LOGGER = LoggerFactory.getLogger(BalancedNumSegmentAssignmentStrategy.class);
@@ -47,22 +44,22 @@ public class BalancedNumSegmentAssignmentStrategy implements SegmentAssignmentSt
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig) {
     SegmentsValidationAndRetentionConfig validationAndRetentionConfig = tableConfig.getValidationConfig();
-    Preconditions.checkState(validationAndRetentionConfig != null, "Validation Config is null");
+    Preconditions.checkState(validationAndRetentionConfig != null, "segmentsConfig is null");
     _replication = tableConfig.getReplication();
     LOGGER.info("Initialized BalancedNumSegmentAssignmentStrategy for table: {} with replication: {}",
         tableConfig.getTableName(), _replication);
   }
 
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
     validateSegmentAssignmentStrategy(instancePartitions);
     return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions, _replication);
   }
 
   @Override
   public Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
     validateSegmentAssignmentStrategy(instancePartitions);
     List<String> instances =
         SegmentAssignmentUtils.getInstancesForNonReplicaGroupBasedAssignment(instancePartitions, _replication);
@@ -74,8 +71,8 @@ private void validateSegmentAssignmentStrategy(InstancePartitions instancePartit
     int numPartitions = instancePartitions.getNumPartitions();
     // Non-replica-group based assignment should have numReplicaGroups and numPartitions = 1
     Preconditions.checkState(numReplicaGroups == 1,
-        "Replica groups should be 1 in order to use BalanceNumSegmentAssignmentStrategy");
+        "Number of replica groups should be 1 in order to use BalancedNumSegmentAssignmentStrategy");
     Preconditions.checkState(numPartitions == 1,
-        "Replica groups should be 1 in order to use BalanceNumSegmentAssignmentStrategy");
+        "Number of instance partitions should be 1 in order to use BalancedNumSegmentAssignmentStrategy");
   }
 }

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java

@@ -33,7 +33,6 @@
 import org.apache.pinot.segment.local.utils.TableConfigUtils;
 import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
-import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -45,17 +44,15 @@ class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy
   private String _tableName;
   private String _partitionColumn;
   private int _replication;
-  private TableConfig _tableConfig;
 
   @Override
   public void init(HelixManager helixManager, TableConfig tableConfig) {
     _helixManager = helixManager;
-    _tableConfig = tableConfig;
     _tableName = tableConfig.getTableName();
     SegmentsValidationAndRetentionConfig validationAndRetentionConfig = tableConfig.getValidationConfig();
-    Preconditions.checkState(validationAndRetentionConfig != null, "Validation Config is null");
+    Preconditions.checkState(validationAndRetentionConfig != null, "segmentsConfig is null");
     _replication = tableConfig.getReplication();
-    _partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
+    _partitionColumn = TableConfigUtils.getPartitionColumn(tableConfig);
     if (_partitionColumn == null) {
       LOGGER.info("Initialized ReplicaGroupSegmentAssignmentStrategy "
           + "with replication: {} without partition column for table: {} ", _replication, _tableName);
@@ -70,7 +67,7 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
    */
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
     checkReplication(instancePartitions, _replication, _tableName);
     int numPartitions = instancePartitions.getNumPartitions();
     int partitionId;
@@ -86,7 +83,7 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
 
   @Override
   public Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+      InstancePartitions instancePartitions) {
     Map<String, Map<String, String>> newAssignment;
     int numPartitions = instancePartitions.getNumPartitions();
 

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategy.java

@@ -23,7 +23,6 @@
 import org.apache.helix.HelixManager;
 import org.apache.pinot.common.assignment.InstancePartitions;
 import org.apache.pinot.spi.config.table.TableConfig;
-import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
 
 
 /**
@@ -40,28 +39,27 @@ public interface SegmentAssignmentStrategy {
   void init(HelixManager helixManager, TableConfig tableConfig);
 
   /**
-   * Assigns segment to instances. The assignment strategy will be configured in
-   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
-   * assignment strategy, this function will be called to assign a new segment
+   * Assigns a segment to instances. The assignment strategy will be configured in the OfflineSegmentAssignment and
+   * RealtimeSegmentAssignment classes, and depending on the type of assignment strategy, this method will be called
+   * to assign a new segment.
    *
    * @param segmentName Name of the segment to be assigned
    * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
    * @param instancePartitions Instance partitions
    * @return List of instances to assign the segment to
    */
   List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType);
+      InstancePartitions instancePartitions);
 
   /**
-   * Re-assigns segment to instances. The assignment strategy will be configured in
-   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
-   * assignment strategy, this function will be called to re-assign a segment
-   * when the InstancePartitions has been changed.
+   * Re-assigns segments to instances. The assignment strategy will be configured in the OfflineSegmentAssignment and
+   * RealtimeSegmentAssignment classes, and depending on the type of assignment strategy, this method will be called
+   * to re-assign segments when the InstancePartitions has been changed.
    *
    * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
    * @param instancePartitions Instance partitions
-   * @return Rebalanced assignment for the segments per assignment strategy
+   * @return Rebalanced assignment for the segments as per the assignment strategy
    */
   Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType);
+      InstancePartitions instancePartitions);
 }