somandal commented on code in PR #15817: URL: https://github.com/apache/pinot/pull/15817#discussion_r2122228140
########## pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java: ########## @@ -1757,67 +1767,26 @@ private static class PartitionIdFetcherImpl implements PartitionIdFetcher { private final String _partitionColumn; private final HelixManager _helixManager; private final boolean _isStrictRealtimeSegmentAssignment; - private final Map<InstancePartitionsType, InstancePartitions> _instancePartitionsMap; private PartitionIdFetcherImpl(String tableNameWithType, @Nullable String partitionColumn, - HelixManager helixManager, boolean isStrictRealtimeSegmentAssignment, - Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) { + HelixManager helixManager, boolean isStrictRealtimeSegmentAssignment) { _tableNameWithType = tableNameWithType; _partitionColumn = partitionColumn; _helixManager = helixManager; _isStrictRealtimeSegmentAssignment = isStrictRealtimeSegmentAssignment; - _instancePartitionsMap = instancePartitionsMap; } @Override public int fetch(String segmentName, boolean isConsuming) { - Integer partitionId; - if (_isStrictRealtimeSegmentAssignment) { - // This is how partitionId is calculated for StrictRealtimeSegmentAssignment. Here partitionId is mandatory - partitionId = - SegmentUtils.getRealtimeSegmentPartitionId(segmentName, _tableNameWithType, _helixManager, - _partitionColumn); - Preconditions.checkState(partitionId != null, "Failed to find partition id for segment: %s of table: %s", - segmentName, _tableNameWithType); - } else { - TableType tableType = TableNameBuilder.getTableTypeFromTableName(_tableNameWithType); - if (tableType == TableType.OFFLINE) { - InstancePartitions instancePartitions = _instancePartitionsMap.get(InstancePartitionsType.OFFLINE); - assert instancePartitions != null; - if (_partitionColumn == null || instancePartitions.getNumPartitions() == 1) { - // Fallback to partitionId 0, in this case batching will not be possible so we will fall back to a full - // rebalance without batching - partitionId = 0; - } else { - // This is how partitionId is calculated for OFFLINE tables - partitionId = SegmentAssignmentUtils.getOfflineSegmentPartitionId(segmentName, _tableNameWithType, - _helixManager, _partitionColumn); - } - } else { - if (isConsuming || !_instancePartitionsMap.containsKey(InstancePartitionsType.COMPLETED)) { - // This is how partitionId is calculated for CONSUMING segments and ONLINE segments without COMPLETED - // instance partitions in RealtimeSegmentAssignment - partitionId = SegmentAssignmentUtils.getRealtimeSegmentPartitionId(segmentName, _tableNameWithType, - _helixManager, _partitionColumn); - } else { - // This is how partitionId is calculated for ONLINE segments when COMPLETED instance partitions exist - // in RealtimeSegmentAssignment - InstancePartitions instancePartitions = _instancePartitionsMap.get(InstancePartitionsType.COMPLETED); - assert instancePartitions != null; - if (_partitionColumn == null || instancePartitions.getNumPartitions() == 1) { - // Fallback to partitionId 0, in this case batching will not be possible so we will fall back to a full - // rebalance without batching - partitionId = 0; - } else { - // This is how partitionId is calculated for REALTIME tables if a partition column exists and if the - // COMPLETED instance partitions has more than 1 partition - partitionId = SegmentAssignmentUtils.getRealtimeSegmentPartitionId(segmentName, _tableNameWithType, - _helixManager, _partitionColumn); - } - } - } + Integer partitionId = Review Comment: discussed offline that it is okay to keep the `InstancePartitions.numPartitions` check decoupled from rebalance perspective. If someone is using strict replica group routing with just 1 partition or with partition column as null, that is not really the correct way to use this (and in the future we may address this via validations etc) I will do some rebalance testing with the changes and get back with findings -- 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