J-HowHuang commented on code in PR #15110: URL: https://github.com/apache/pinot/pull/15110#discussion_r1972259795
########## pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java: ########## @@ -886,6 +886,188 @@ public void testRebalanceWithTiersAndInstanceAssignments() _helixResourceManager.deleteOfflineTable(TIERED_TABLE_NAME); } + + @Test + public void testRebalanceWithMinimizeDataMovementInstanceAssignments() + throws Exception { + // TODO: try balanced (default) instance assignment + 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 + numServers, 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); Review Comment: a299268 ########## pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java: ########## @@ -886,6 +886,188 @@ public void testRebalanceWithTiersAndInstanceAssignments() _helixResourceManager.deleteOfflineTable(TIERED_TABLE_NAME); } + + @Test + public void testRebalanceWithMinimizeDataMovementInstanceAssignments() + throws Exception { + // TODO: try balanced (default) instance assignment + 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, Review Comment: a299268 -- 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