klsince commented on code in PR #16017:
URL: https://github.com/apache/pinot/pull/16017#discussion_r2165071374
##########
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:
correct. good catch about the input param. Not introduced in this PR, but
that map should include a COMPLETED entry to validate the sanity checks. Will
fix that.
Will add tests on rebalanceTable() too, to validate the sanity checks on
COMPLETED tag.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
