Address feedback

ankitsultana · ankitsultana · commit a9ab0738ddd0 · 2022-08-10T10:34:12.000-04:00
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
@@ -1701,11 +1701,12 @@ private void assignInstances(TableConfig tableConfig, boolean override) {
           LOGGER.info("Persisting instance partitions: {}", instancePartitions);
           InstancePartitionsUtils.persistInstancePartitions(_propertyStore, instancePartitions);
         } else {
+          String referenceInstancePartitionsName =
+              tableConfig.getInstancePartitionsMap().get(instancePartitionsType);
           instancePartitions = InstancePartitionsUtils.fetchInstancePartitionsWithRename(_propertyStore,
-              tableConfig.getInstancePartitionsMap().get(instancePartitionsType),
-              instancePartitionsType.getInstancePartitionsName(rawTableName));
+              referenceInstancePartitionsName, instancePartitionsType.getInstancePartitionsName(rawTableName));
           LOGGER.info("Persisting instance partitions: {} (referencing {})", instancePartitions,
-              tableConfig.getInstancePartitionsMap().get(instancePartitionsType));
+              referenceInstancePartitionsName);
           InstancePartitionsUtils.persistInstancePartitions(_propertyStore, instancePartitions);
         }
       }
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java
@@ -424,12 +424,13 @@ private InstancePartitions getInstancePartitions(TableConfig tableConfig,
         boolean hasPreConfiguredInstancePartitions = TableConfigUtils.hasPreConfiguredInstancePartitions(tableConfig,
             instancePartitionsType);
         if (hasPreConfiguredInstancePartitions) {
+          String referenceInstancePartitionsName = tableConfig.getInstancePartitionsMap().get(instancePartitionsType);
           InstancePartitions instancePartitions = InstancePartitionsUtils.fetchInstancePartitionsWithRename(
-              _helixManager.getHelixPropertyStore(), tableConfig.getInstancePartitionsMap().get(instancePartitionsType),
+              _helixManager.getHelixPropertyStore(), referenceInstancePartitionsName,
               instancePartitionsType.getInstancePartitionsName(rawTableName));
           if (!dryRun) {
             LOGGER.info("Persisting instance partitions: {} (referencing {})", instancePartitions,
-                tableConfig.getInstancePartitionsMap().get(instancePartitionsType));
+                referenceInstancePartitionsName);
             InstancePartitionsUtils.persistInstancePartitions(_helixManager.getHelixPropertyStore(),
                 instancePartitions);
           }
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java
@@ -85,6 +85,10 @@ public class TableConfig extends BaseJsonConfig {
   private RoutingConfig _routingConfig;
   private QueryConfig _queryConfig;
   private Map<InstancePartitionsType, InstanceAssignmentConfig> _instanceAssignmentConfigMap;
+
+  @JsonPropertyDescription(value = "Point to an existing instance partitions")
+  private Map<InstancePartitionsType, String> _instancePartitionsMap;
+
   private List<FieldConfig> _fieldConfigList;
 
   @JsonPropertyDescription(value = "upsert related config")
@@ -102,9 +106,6 @@ public class TableConfig extends BaseJsonConfig {
   @JsonPropertyDescription(value = "Configs for Table config tuner")
   private List<TunerConfig> _tunerConfigList;
 
-  @JsonPropertyDescription(value = "Point to an existing instance partitions")
-  private Map<InstancePartitionsType, String> _instancePartitionsMap;
-
   @JsonCreator
   public TableConfig(@JsonProperty(value = TABLE_NAME_KEY, required = true) String tableName,
       @JsonProperty(value = TABLE_TYPE_KEY, required = true) String tableType,
@@ -261,6 +262,15 @@ public void setInstanceAssignmentConfigMap(
     _instanceAssignmentConfigMap = instanceAssignmentConfigMap;
   }
 
+  @JsonProperty(INSTANCE_PARTITIONS_MAP_CONFIG_KEY)
+  public Map<InstancePartitionsType, String> getInstancePartitionsMap() {
+    return _instancePartitionsMap;
+  }
+
+  public void setInstancePartitionsMap(Map<InstancePartitionsType, String> instancePartitionsMap) {
+    _instancePartitionsMap = instancePartitionsMap;
+  }
+
   @JsonProperty(FIELD_CONFIG_LIST_KEY)
   @Nullable
   public List<FieldConfig> getFieldConfigList() {
@@ -328,13 +338,4 @@ public List<TunerConfig> getTunerConfigsList() {
   public void setTunerConfigsList(List<TunerConfig> tunerConfigList) {
     _tunerConfigList = tunerConfigList;
   }
-
-  @JsonProperty(INSTANCE_PARTITIONS_MAP_CONFIG_KEY)
-  public Map<InstancePartitionsType, String> getInstancePartitionsMap() {
-    return _instancePartitionsMap;
-  }
-
-  public void setInstancePartitionsMap(Map<InstancePartitionsType, String> instancePartitionsMap) {
-    _instancePartitionsMap = instancePartitionsMap;
-  }
 }
