Fix condition check while fetching instance partitions (apache#14910)

* use table level ip * add unit tests * update function name * add file header * rename function * fix checkstyle error * address comments * Update pinot-common/src/main/java/org/apache/pinot/common/assignment/InstancePartitionsUtils.java Co-authored-by: Jia Guo <[email protected]> * Fix method signature in InstancePartitionsUtils.java * Retrigger Tests --------- Co-authored-by: Endong Zhu <[email protected]> Co-authored-by: Jia Guo <[email protected]>
xiangfu0 · Jan 29, 2025 · 156be80 · 156be80
1 parent bfe5728
commit 156be80
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 3 deletions.
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
@@ -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<ZNRecord> 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/...-common/src/test/java/org/apache/pinot/common/assignment/InstancePartitionsUtilsTest.java b/...-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<InstancePartitionsType, String> instancePartitionsMap = new HashMap<>();
+    instancePartitionsMap.put(InstancePartitionsType.OFFLINE, "testPartitions");
+    Map<String, InstanceAssignmentConfig> 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<InstancePartitionsType, String> instancePartitionsMap = new HashMap<>();
+    instancePartitionsMap.put(InstancePartitionsType.OFFLINE, "testPartitions");
+    Map<String, InstanceAssignmentConfig> 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<String> 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);
+  }
+}