somandal commented on code in PR #16017:
URL: https://github.com/apache/pinot/pull/16017#discussion_r2136612615
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java:
##########
@@ -255,9 +279,74 @@ public void testAssignSegmentWithOfflineSegment() {
}
}
- @Test(expectedExceptions = IllegalStateException.class)
- public void testAssignSegmentToCompletedServers() {
- _segmentAssignment.assignSegment("seg01", new TreeMap<>(), new
TreeMap<>());
+ @Test
+ public void testRebalanceDedupTableWithTiers() {
+ SegmentAssignment segmentAssignment = createSegmentAssignment("dedup");
+ Map<InstancePartitionsType, InstancePartitions>
onlyConsumingInstancePartitionMap =
+ ImmutableMap.of(InstancePartitionsType.CONSUMING,
_instancePartitionsMap.get(InstancePartitionsType.CONSUMING));
+ Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
+ Set<String> segmentsOnTier = new HashSet<>();
+ for (int segmentId = 0; segmentId < 6; segmentId++) {
+ String segmentName = _segments.get(segmentId);
+ if (segmentId < 3) {
+ segmentsOnTier.add(segmentName);
+ }
+ List<String> instancesAssigned =
+ segmentAssignment.assignSegment(segmentName, currentAssignment,
_instancePartitionsMap);
+ currentAssignment.put(segmentName,
+ SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned,
SegmentStateModel.ONLINE));
+ }
+ String tierName = "coldTier";
+ List<Tier> sortedTiers = createSortedTiers(tierName, segmentsOnTier);
+ Map<String, InstancePartitions> tierInstancePartitionsMap =
createTierInstancePartitionsMap(tierName, 3);
+ RebalanceConfig rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setIncludeConsuming(true);
+ Map<String, Map<String, String>> newAssignment =
+ segmentAssignment.rebalanceTable(currentAssignment,
onlyConsumingInstancePartitionMap, sortedTiers,
+ tierInstancePartitionsMap, rebalanceConfig);
+ assertEquals(newAssignment.size(), currentAssignment.size());
+ for (String segName : currentAssignment.keySet()) {
+ if (segmentsOnTier.contains(segName)) {
+ assertTrue(newAssignment.get(segName).keySet().stream().allMatch(s ->
s.startsWith(tierName)));
+ } else {
+ assertTrue(
+ newAssignment.get(segName).keySet().stream().allMatch(s ->
s.startsWith(CONSUMING_INSTANCE_NAME_PREFIX)));
+ }
+ }
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class,
expectedExceptionsMessageRegExp = "Tiers must not be "
+ + "specified for table.*")
+ public void testRebalanceUpsertTableWithTiers() {
+ SegmentAssignment segmentAssignment = createSegmentAssignment("upsert");
+ Map<InstancePartitionsType, InstancePartitions>
onlyConsumingInstancePartitionMap =
+ ImmutableMap.of(InstancePartitionsType.CONSUMING,
_instancePartitionsMap.get(InstancePartitionsType.CONSUMING));
+ Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
+ Set<String> segmentsOnTier = new HashSet<>();
+ for (int segmentId = 0; segmentId < 6; segmentId++) {
+ String segmentName = _segments.get(segmentId);
+ if (segmentId < 3) {
+ segmentsOnTier.add(segmentName);
+ }
+ List<String> instancesAssigned =
+ segmentAssignment.assignSegment(segmentName, currentAssignment,
_instancePartitionsMap);
+ currentAssignment.put(segmentName,
+ SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned,
SegmentStateModel.ONLINE));
+ }
+
+ String tierName = "coldTier";
+ List<Tier> sortedTiers = createSortedTiers(tierName, segmentsOnTier);
+ Map<String, InstancePartitions> tierInstancePartitionsMap =
createTierInstancePartitionsMap(tierName, 3);
+ RebalanceConfig rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setIncludeConsuming(true);
+ segmentAssignment.rebalanceTable(currentAssignment,
onlyConsumingInstancePartitionMap, sortedTiers,
+ tierInstancePartitionsMap, rebalanceConfig);
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class, dataProvider =
"tableTypes")
+ public void testAssignSegmentToCompletedServers(String tableType) {
+ SegmentAssignment segmentAssignment = createSegmentAssignment(tableType);
+ segmentAssignment.assignSegment("seg01", new TreeMap<>(), new TreeMap<>());
Review Comment:
is this test ensuring that we don't move segments to "COMPLETED" instance
partitions? but the passed in instancePartitionsMap is empty?
Can we also have a test to reject if instancePartitionsMap has COMPLETED?
perhaps also have similar test for rebalanceTables() API?
--
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
