diff --git a/pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java index 759d387af47a..47b0319f0382 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java @@ -60,12 +60,12 @@ public static String getInstancePartitionsName(String tableName, String instance * Fetches the instance partitions from Helix property store if it exists, or computes it for backward-compatibility. */ public static InstancePartitions fetchOrComputeInstancePartitions(HelixManager helixManager, TableConfig tableConfig, - InstancePartitionsType instancePartitionsType) { +InstancePartitionsType instancePartitionsType) { String tableNameWithType = tableConfig.getTableName(); String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); - // If table has pre-configured instance partitions. - if (TableConfigUtils.hasPreConfiguredInstancePartitions(tableConfig, instancePartitionsType)) { + // If table has pre-configured table-level instance partitions + if (shouldFetchPreConfiguredInstancePartitions(tableConfig, instancePartitionsType)) { return fetchInstancePartitionsWithRename(helixManager.getHelixPropertyStore(), tableConfig.getInstancePartitionsMap().get(instancePartitionsType), instancePartitionsType.getInstancePartitionsName(rawTableName)); @@ -193,4 +193,10 @@ public static void removeTierInstancePartitions(HelixPropertyStore pro removeInstancePartitions(propertyStore, instancePartition.getInstancePartitionsName()); }); } + + public static boolean shouldFetchPreConfiguredInstancePartitions(TableConfig tableConfig, + InstancePartitionsType instancePartitionsType) { + return TableConfigUtils.hasPreConfiguredInstancePartitions(tableConfig, instancePartitionsType) + && !InstanceAssignmentConfigUtils.isMirrorServerSetAssignment(tableConfig, instancePartitionsType); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/assignment/InstancePartitionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/assignment/InstancePartitionsUtilsTest.java new file mode 100644 index 000000000000..bf80bf787b01 --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/assignment/InstancePartitionsUtilsTest.java @@ -0,0 +1,84 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.common.assignment; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.assignment.InstanceAssignmentConfig; +import org.apache.pinot.spi.config.table.assignment.InstanceConstraintConfig; +import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType; +import org.apache.pinot.spi.config.table.assignment.InstanceReplicaGroupPartitionConfig; +import org.apache.pinot.spi.config.table.assignment.InstanceTagPoolConfig; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class InstancePartitionsUtilsTest { + + @Test + public void testShouldFetchPreConfiguredInstancePartitions() { + Map instancePartitionsMap = new HashMap<>(); + instancePartitionsMap.put(InstancePartitionsType.OFFLINE, "testPartitions"); + Map instanceAssignmentConfigMap = new HashMap<>(); + instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(), + getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR)); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable") + .setInstancePartitionsMap(instancePartitionsMap) + .setInstanceAssignmentConfigMap(instanceAssignmentConfigMap) + .build(); + + Assert.assertTrue(InstancePartitionsUtils.shouldFetchPreConfiguredInstancePartitions(tableConfig, + InstancePartitionsType.OFFLINE)); + } + + @Test + public void testShouldFetchPreConfiguredInstancePartitionsMirrorServerSet() { + Map instancePartitionsMap = new HashMap<>(); + instancePartitionsMap.put(InstancePartitionsType.OFFLINE, "testPartitions"); + Map instanceAssignmentConfigMap = new HashMap<>(); + instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(), + getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.MIRROR_SERVER_SET_PARTITION_SELECTOR)); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable") + .setInstancePartitionsMap(instancePartitionsMap) + .setInstanceAssignmentConfigMap(instanceAssignmentConfigMap) + .build(); + + Assert.assertFalse(InstancePartitionsUtils.shouldFetchPreConfiguredInstancePartitions(tableConfig, + InstancePartitionsType.OFFLINE)); + } + + private static InstanceAssignmentConfig getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector + partitionSelector) { + InstanceTagPoolConfig instanceTagPoolConfig = + new InstanceTagPoolConfig("tag", true, 1, null); + List constraints = new ArrayList<>(); + constraints.add("constraints1"); + InstanceConstraintConfig instanceConstraintConfig = new InstanceConstraintConfig(constraints); + InstanceReplicaGroupPartitionConfig instanceReplicaGroupPartitionConfig = + new InstanceReplicaGroupPartitionConfig(true, 1, 1, + 1, 1, 1, true, + null); + return new InstanceAssignmentConfig(instanceTagPoolConfig, instanceConstraintConfig, + instanceReplicaGroupPartitionConfig, partitionSelector.name(), false); + } +}