klsince commented on code in PR #10255: URL: https://github.com/apache/pinot/pull/10255#discussion_r1123829006
########## pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotInstanceAssignmentRestletResource.java: ########## @@ -315,16 +407,16 @@ private void removeInstancePartitionsHelper(String instancePartitionsName) { @Path("/tables/{tableName}/replaceInstance") @Authenticate(AccessType.CREATE) @ApiOperation(value = "Replace an instance in the instance partitions") - public Map<InstancePartitionsType, InstancePartitions> replaceInstance( + public Map<String, InstancePartitions> replaceInstance( @ApiParam(value = "Name of the table") @PathParam("tableName") String tableName, - @ApiParam(value = "OFFLINE|CONSUMING|COMPLETED") @QueryParam("type") @Nullable - InstancePartitionsType instancePartitionsType, + @ApiParam(value = "OFFLINE|CONSUMING|COMPLETED|Name of the tier") @QueryParam("type") @Nullable Review Comment: nit: tierName for short (and change a few other places for consistency) ########## pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotInstanceAssignmentRestletResource.java: ########## @@ -207,22 +247,42 @@ public Map<InstancePartitionsType, InstancePartitions> assignInstances( * @param instanceConfigs list of instance configs * @param instancePartitionsType type of instancePartitions */ - private void assignInstancesForInstancePartitionsType( - Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, TableConfig tableConfig, - List<InstanceConfig> instanceConfigs, InstancePartitionsType instancePartitionsType) { + private void assignInstancesForInstancePartitionsType(Map<String, InstancePartitions> instancePartitionsMap, + TableConfig tableConfig, List<InstanceConfig> instanceConfigs, InstancePartitionsType instancePartitionsType) { String tableNameWithType = tableConfig.getTableName(); if (TableConfigUtils.hasPreConfiguredInstancePartitions(tableConfig, instancePartitionsType)) { String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); - instancePartitionsMap.put(instancePartitionsType, InstancePartitionsUtils.fetchInstancePartitionsWithRename( - _resourceManager.getPropertyStore(), tableConfig.getInstancePartitionsMap().get(instancePartitionsType), - instancePartitionsType.getInstancePartitionsName(rawTableName))); + instancePartitionsMap.put(instancePartitionsType.toString(), + InstancePartitionsUtils.fetchInstancePartitionsWithRename(_resourceManager.getPropertyStore(), + tableConfig.getInstancePartitionsMap().get(instancePartitionsType), + instancePartitionsType.getInstancePartitionsName(rawTableName))); return; } InstancePartitions existingInstancePartitions = InstancePartitionsUtils.fetchInstancePartitions(_resourceManager.getHelixZkManager().getHelixPropertyStore(), InstancePartitionsUtils.getInstancePartitionsName(tableNameWithType, instancePartitionsType.toString())); - instancePartitionsMap.put(instancePartitionsType, new InstanceAssignmentDriver(tableConfig) - .assignInstances(instancePartitionsType, instanceConfigs, existingInstancePartitions)); + instancePartitionsMap.put(instancePartitionsType.toString(), + new InstanceAssignmentDriver(tableConfig).assignInstances(instancePartitionsType, instanceConfigs, + existingInstancePartitions)); + } + + private void assignInstancesForTier(Map<String, InstancePartitions> instancePartitionsMap, TableConfig tableConfig, + List<InstanceConfig> instanceConfigs, String tierName) { + if (CollectionUtils.isNotEmpty(tableConfig.getTierConfigsList()) + && tableConfig.getInstanceAssignmentConfigMap() != null) { + for (TierConfig tierConfig : tableConfig.getTierConfigsList()) { + if ((tierConfig.getName().equals(tierName) || tierName == null) Review Comment: Q: when tierName is null, why we still need to assignInstances for tierConfig.getName()? ########## pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java: ########## @@ -524,31 +525,73 @@ private List<Tier> getSortedTiers(TableConfig tableConfig) { } @Nullable - private Map<String, InstancePartitions> getTierToInstancePartitionsMap(String tableNameWithType, - @Nullable List<Tier> sortedTiers) { + private Map<String, InstancePartitions> getTierToInstancePartitionsMap(TableConfig tableConfig, + @Nullable List<Tier> sortedTiers, boolean reassignInstances, boolean bootstrap, boolean dryRun) { 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); - tierToInstancePartitionsMap.put(tier.getName(), getInstancePartitionsForTier(tier, tableNameWithType)); + tableConfig.getTableName()); + tierToInstancePartitionsMap.put(tier.getName(), + getInstancePartitionsForTier(tableConfig, tier, reassignInstances, bootstrap, dryRun)); } return tierToInstancePartitionsMap; } /** - * Creates a default instance assignment for the tier. - * TODO: We only support default server-tag based assignment currently. - * In next iteration, we will add InstanceAssignmentConfig to the TierConfig and also support persisting of the - * InstancePartitions to zk. - * Then we'll be able to support replica group assignment while creating InstancePartitions for tiers + * Computes the instance partitions for the given tier. If table's instanceAssignmentConfigMap has an entry for the + * tier, it's used to calculate the instance partitions. Else default instance partitions are returned */ - private InstancePartitions getInstancePartitionsForTier(Tier tier, String tableNameWithType) { + private InstancePartitions getInstancePartitionsForTier(TableConfig tableConfig, Tier tier, boolean reassignInstances, + boolean bootstrap, boolean dryRun) { PinotServerTierStorage storage = (PinotServerTierStorage) tier.getStorage(); - return InstancePartitionsUtils.computeDefaultInstancePartitionsForTag(_helixManager, tableNameWithType, - tier.getName(), storage.getServerTag()); + InstancePartitions defaultInstancePartitions = + InstancePartitionsUtils.computeDefaultInstancePartitionsForTag(_helixManager, tableConfig.getTableName(), + tier.getName(), storage.getServerTag()); + + if (tableConfig.getInstanceAssignmentConfigMap() == null || !tableConfig.getInstanceAssignmentConfigMap() + .containsKey(tier.getName())) { + LOGGER.info("Skipping fetching/computing instance partitions for tier {} for table: {}", tier.getName(), + tableConfig.getTableName()); + if (!dryRun) { + String instancePartitionsName = + InstancePartitionsUtils.getInstancePartitonNameForTier(tableConfig.getTableName(), tier.getName()); + LOGGER.info("Removing instance partitions: {} from ZK if it exists", instancePartitionsName); + InstancePartitionsUtils.removeInstancePartitions(_helixManager.getHelixPropertyStore(), instancePartitionsName); + } + return defaultInstancePartitions; + } + + String tableNameWithType = tableConfig.getTableName(); + String instancePartitionName = Review Comment: nit: instancePartition`s`Name to be a bit consistent ########## pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java: ########## @@ -1966,6 +1991,7 @@ public void deleteRealtimeTable(String tableName, @Nullable String retentionPeri InstancePartitionsType.CONSUMING.getInstancePartitionsName(rawTableName)); InstancePartitionsUtils.removeInstancePartitions(_propertyStore, InstancePartitionsType.COMPLETED.getInstancePartitionsName(rawTableName)); + InstancePartitionsUtils.removeTierInstancePartitions(_propertyStore, rawTableName); Review Comment: nit: just to be consistent with the method above, move this one below the info log and add info log for deleting tier instance partitions? ########## pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java: ########## @@ -93,6 +94,11 @@ public static InstancePartitions fetchInstancePartitions(HelixPropertyStore<ZNRe return znRecord != null ? InstancePartitions.fromZNRecord(znRecord) : null; } + public static String getInstancePartitonNameForTier(String tableName, String tierName) { Review Comment: nit: the naming convention seems getInstancePartiton`s`NameForTier ########## pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotInstanceAssignmentRestletResource.java: ########## @@ -80,39 +83,68 @@ public class PinotInstanceAssignmentRestletResource { @Produces(MediaType.APPLICATION_JSON) @Path("/tables/{tableName}/instancePartitions") @ApiOperation(value = "Get the instance partitions") - public Map<InstancePartitionsType, InstancePartitions> getInstancePartitions( + public Map<String, InstancePartitions> getInstancePartitions( @ApiParam(value = "Name of the table") @PathParam("tableName") String tableName, - @ApiParam(value = "OFFLINE|CONSUMING|COMPLETED") @QueryParam("type") @Nullable - InstancePartitionsType instancePartitionsType) { - Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap = new TreeMap<>(); + @ApiParam(value = "OFFLINE|CONSUMING|COMPLETED|Name of the tier") @QueryParam("type") @Nullable String type) { + Map<String, InstancePartitions> instancePartitionsMap = new TreeMap<>(); String rawTableName = TableNameBuilder.extractRawTableName(tableName); TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableName); if (tableType != TableType.REALTIME) { - if (instancePartitionsType == InstancePartitionsType.OFFLINE || instancePartitionsType == null) { - InstancePartitions offlineInstancePartitions = InstancePartitionsUtils - .fetchInstancePartitions(_resourceManager.getPropertyStore(), + if (InstancePartitionsType.OFFLINE.toString().equals(type) || type == null) { + InstancePartitions offlineInstancePartitions = + InstancePartitionsUtils.fetchInstancePartitions(_resourceManager.getPropertyStore(), InstancePartitionsType.OFFLINE.getInstancePartitionsName(rawTableName)); if (offlineInstancePartitions != null) { - instancePartitionsMap.put(InstancePartitionsType.OFFLINE, offlineInstancePartitions); + instancePartitionsMap.put(InstancePartitionsType.OFFLINE.toString(), offlineInstancePartitions); } } } if (tableType != TableType.OFFLINE) { - if (instancePartitionsType == InstancePartitionsType.CONSUMING || instancePartitionsType == null) { - InstancePartitions consumingInstancePartitions = InstancePartitionsUtils - .fetchInstancePartitions(_resourceManager.getPropertyStore(), + if (InstancePartitionsType.CONSUMING.toString().equals(type) || type == null) { + InstancePartitions consumingInstancePartitions = + InstancePartitionsUtils.fetchInstancePartitions(_resourceManager.getPropertyStore(), InstancePartitionsType.CONSUMING.getInstancePartitionsName(rawTableName)); if (consumingInstancePartitions != null) { - instancePartitionsMap.put(InstancePartitionsType.CONSUMING, consumingInstancePartitions); + instancePartitionsMap.put(InstancePartitionsType.CONSUMING.toString(), consumingInstancePartitions); } } - if (instancePartitionsType == InstancePartitionsType.COMPLETED || instancePartitionsType == null) { - InstancePartitions completedInstancePartitions = InstancePartitionsUtils - .fetchInstancePartitions(_resourceManager.getPropertyStore(), + if (InstancePartitionsType.COMPLETED.toString().equals(type) || type == null) { + InstancePartitions completedInstancePartitions = + InstancePartitionsUtils.fetchInstancePartitions(_resourceManager.getPropertyStore(), InstancePartitionsType.COMPLETED.getInstancePartitionsName(rawTableName)); if (completedInstancePartitions != null) { - instancePartitionsMap.put(InstancePartitionsType.COMPLETED, completedInstancePartitions); + instancePartitionsMap.put(InstancePartitionsType.COMPLETED.toString(), completedInstancePartitions); + } + } + } + Review Comment: use ` List<TableConfig> tableConfigs = Arrays.asList(_resourceManager.getRealtimeTableConfig(tableName), _resourceManager.getOfflineTableConfig(tableName));` I seen used in another method below, to combine the two big if-blocks? -- 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