somandal commented on code in PR #15110:
URL: https://github.com/apache/pinot/pull/15110#discussion_r1972440509
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java:
##########
@@ -886,6 +886,226 @@ public void testRebalanceWithTiersAndInstanceAssignments()
_helixResourceManager.deleteOfflineTable(TIERED_TABLE_NAME);
}
+ @Test
+ public void testRebalanceWithMinimizeDataMovementBalanced() throws Exception
{
+ int numServers = 6;
+ for (int i = 0; i < numServers; i++) {
+ addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX +
i, true);
+ }
+
+ // Create the table with default balanced segment assignment
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build();
+
+ addDummySchema(RAW_TABLE_NAME);
+ _helixResourceManager.addTable(tableConfig);
+
+ // Add the segments
+ int numSegments = 10;
+ long nowInDays = TimeUnit.MILLISECONDS.toDays(System.currentTimeMillis());
+
+ for (int i = 0; i < numSegments; i++) {
+ _helixResourceManager.addNewSegment(OFFLINE_TIERED_TABLE_NAME,
+
SegmentMetadataMockUtils.mockSegmentMetadataWithEndTimeInfo(TIERED_TABLE_NAME,
SEGMENT_NAME_PREFIX + i,
+ nowInDays), null);
+ }
+
+ Map<String, Map<String, String>> oldSegmentAssignment =
+
_helixResourceManager.getTableIdealState(OFFLINE_TIERED_TABLE_NAME).getRecord().getMapFields();
+
+ TableRebalancer tableRebalancer = new TableRebalancer(_helixManager);
+
+ // Try dry-run summary mode
+ RebalanceConfig rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setDryRun(true);
+ rebalanceConfig.setSummary(true);
+ rebalanceConfig.setReassignInstances(true);
+ rebalanceConfig.setMinimizeDataMovement(true);
+ RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig,
rebalanceConfig, null);
+
+ RebalanceSummaryResult rebalanceSummaryResult =
rebalanceResult.getRebalanceSummaryResult();
+ assertNotNull(rebalanceSummaryResult);
+ assertNotNull(rebalanceSummaryResult.getServerInfo());
+ RebalanceSummaryResult.ServerInfo rebalanceServerInfo =
rebalanceSummaryResult.getServerInfo();
+ assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP);
+
assertEquals(rebalanceServerInfo.getNumServers().getExpectedValueAfterRebalance(),
numServers);
+
+ rebalanceResult = tableRebalancer.rebalance(tableConfig, new
RebalanceConfig(), null);
+ assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP);
+ // Segment assignment should not change
+ assertEquals(rebalanceResult.getSegmentAssignment(), oldSegmentAssignment);
+
+ // add one server instance
+ addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX,
true);
+
+ // Table without instance assignment config should work fine (ignore) with
the minimizeDataMovement flag set
+ // Try dry-run summary mode
+ rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setDryRun(true);
+ rebalanceConfig.setSummary(true);
+ rebalanceConfig.setReassignInstances(true);
+ rebalanceConfig.setMinimizeDataMovement(true);
+ rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig,
null);
+
+ rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult();
+ assertNotNull(rebalanceSummaryResult);
+ assertNotNull(rebalanceSummaryResult.getServerInfo());
+ rebalanceServerInfo = rebalanceSummaryResult.getServerInfo();
+ // Should see the added server
+ assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE);
+
assertEquals(rebalanceServerInfo.getNumServers().getValueBeforeRebalance(),
numServers);
+
assertEquals(rebalanceServerInfo.getNumServers().getExpectedValueAfterRebalance(),
numServers + 1);
+
+
+ // Check if the instance assignment is the same as the one without
minimizeDataMovement flag set
+ rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setDryRun(true);
+ rebalanceConfig.setSummary(true);
+ rebalanceConfig.setReassignInstances(true);
+ rebalanceConfig.setMinimizeDataMovement(false);
+ RebalanceResult rebalanceResultWithoutMinimized =
tableRebalancer.rebalance(tableConfig, rebalanceConfig, null);
+
+ assertEquals(rebalanceResult.getInstanceAssignment(),
rebalanceResultWithoutMinimized.getInstanceAssignment());
+
+ // Rebalance
+ rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setReassignInstances(true);
+ rebalanceConfig.setMinimizeDataMovement(true);
+ rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig,
null);
+ // Should see the added server in the instance assignment
+ assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE);
+
assertEquals(rebalanceResult.getInstanceAssignment().get(InstancePartitionsType.OFFLINE).getInstances(0,
0).size(), numServers + 1);
+
+ }
+
+ @Test
+ public void testRebalanceWithMinimizeDataMovementInstanceAssignments()
+ throws Exception {
+ int numServers = 6;
+ for (int i = 0; i < numServers; i++) {
+ addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX +
i, true);
+ }
+
+ // One instance per replica group, no partition
+ InstanceAssignmentConfig instanceAssignmentConfig = new
InstanceAssignmentConfig(
+ new InstanceTagPoolConfig(TagNameUtils.getOfflineTagForTenant(null),
false, 0, null), null,
+ new InstanceReplicaGroupPartitionConfig(true, 0, NUM_REPLICAS, 1, 0,
0, false,
+ null), null, false);
+
+ // Create the table with default balanced segment assignment
Review Comment:
nit: is this comment correct? thought this test is for replicaGroups?
--
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
