You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by GitBox <gi...@apache.org> on 2022/08/31 20:52:05 UTC

[GitHub] [pinot] GSharayu opened a new pull request, #9309: Introduce Segment AssignmentStrategy Interface

GSharayu opened a new pull request, #9309:
URL: https://github.com/apache/pinot/pull/9309

   Currently, Offline/RealtimeSegmentAssignment covers all the available segment assignment strategies. Single implementation works for both replica group and non replica group based tables.A lot of concepts are coupled together in the existing segment assignment implementation. (Tiered storage, Realtime’s consuming/completed concept etc)
   
   Assignment strategies (e.g. balanced, replica-group etc) needs to be decoupled from SegmentAssignment.
   
   Design strategy doc:
   https://docs.google.com/document/d/16sGvC8UPHcaYJvUXHB2wODsaeNYkXj4lAhaJiSk9Jiw/edit?usp=sharing
   
   Strategy Interface for real time consuming segments will be a follow up PR
   #9047 
   #8585 
   cc @snleee 


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r961188711


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized OfflineBalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableName, _tableName, _replication);

Review Comment:
   ```suggestion
               + "{} with replication: {}", _tableName, _replication);
   ```



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970070945


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,109 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, InstancePartitions instancePartitions) {

Review Comment:
   So the idea here will be to fetch the instancePartitions for given InstancePartitionType, so we dont particularly need instanceType even in consuming segments



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969865584


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   @GSharayu yeah sorry, after looking at this again, I'm not sure if I see a way to remove `_tableConfig`. I thought by updated you meant you've removed it, so therefore my nit comment yesterday. Can @snleee elaborate if removing `_tableConfig` is still feasible for the current setup? 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] klsince commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
klsince commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r1018496195


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   hi @GSharayu , I'm testing locally with RealtimeQuickStart customized a little bit, and I came into NPE issue like below. Basically, in my local test setup, the tableConfig has some tierConfigs; and SegmentRelocator is enabled to run periodically to move segments around based on data age as set in tierConfigs. 
   
   As suggested above, I feel it'd be safer to init segmentAssignmentStrategy with a default policy. 
   
   ```
   2022/11/09 14:43:09.474 INFO [RealtimeSegmentAssignment] [async-task-thread-1] Rebalancing tier: myTier for table: airlineStats_REALTIME with bootstrap: false, instance partitions: {"instancePartitionsName":"airlineStats
   _myTier","partitionToInstancesMap":{"0_0":["Server_192.168.0.106_7000"]}}
   2022/11/09 14:43:09.475 WARN [TableRebalancer] [async-task-thread-1] Caught exception while calculating target assignment for table: airlineStats_REALTIME, aborting the rebalance
   java.lang.NullPointerException: null
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.reassignSegments(BaseSegmentAssignment.java:158) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.rebalanceTiers(BaseSegmentAssignment.java:129) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.RealtimeSegmentAssignment.rebalanceTable(RealtimeSegmentAssignment.java:201) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.rebalance.TableRebalancer.rebalance(TableRebalancer.java:219) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.relocation.SegmentRelocator.lambda$processTable$0(SegmentRelocator.java:137) ~[classes/:?]
           at java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:515) [?:?]
           at java.util.concurrent.FutureTask.run(FutureTask.java:264) [?:?]
           at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128) [?:?]
           at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628) [?:?]
           at java.lang.Thread.run(Thread.java:829) [?:?]
   ```
   
   cc @npawar for awareness ^



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r971256087


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -90,9 +89,20 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
     InstancePartitions instancePartitions = typeToInstancePartitions.getValue();
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
+
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -53,62 +49,57 @@ protected int getReplication(TableConfig tableConfig) {
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitions, InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] jackjlli commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
jackjlli commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970187650


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    _assignmentStrategyMap = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitionsMap);

Review Comment:
   According the caller method name, it's a setter method; while in this line, it's a getter method. It'd be good to come up with a better name in order to align with the logic we're doing here.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r968719287


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   In the real time tables, we would want to get assignment strategies for both CONSUMING and COMPLETED segments at one time, so the idea was to keep a map. If we change this to fetch single assignment strategy, then probably, we would want to call same function for 2 different instance partition types, instead we can populate map and use the key to determine strategy.
   
   cc @snleee 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r976041452


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,87 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {

Review Comment:
   Let's rename this to be  `Balanced` to be sync with the documentation
   
   https://docs.pinot.apache.org/operators/operating-pinot/segment-assignment



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r964213302


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -185,7 +157,8 @@ protected Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, S
    * Rebalances segments in the current assignment using the instancePartitions and returns new assignment
    */
   protected Map<String, Map<String, String>> reassignSegments(String instancePartitionType,
-      Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, boolean bootstrap) {
+      Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, boolean bootstrap,
+      SegmentAssignmentStrategy segmentAssignmentStrategy, InstancePartitionsType instancePartitionsType) {

Review Comment:
   Seems a bit redundant to have both `String instancePartitionType` and `InstancePartitionsType instancePartitionsType`, wondering if there's a way to collapse these (however, seems like the first `String instancePartitionType` is mixed with tier storage concept from before your change). 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] Jackie-Jiang commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
Jackie-Jiang commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r977053088


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,107 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy

Review Comment:
   (code style) Can you please reformat all the changes with [Pinot Style](https://docs.pinot.apache.org/developers/developers-and-contributors/code-setup#setup-ide)? I can see several places not following the format convention



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,87 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableNameWithType;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableNameWithType = tableConfig.getTableName();
+    SegmentsValidationAndRetentionConfig validationAndRetentionConfig = tableConfig.getValidationConfig();
+    Preconditions.checkState(validationAndRetentionConfig != null, "Validation Config is null");
+    _replication = validationAndRetentionConfig.getReplicationNumber();
+    LOGGER.info("Initialized BalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableNameWithType, _replication);

Review Comment:
   (minor) Don't break the string into 2 parts unless it is too long



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,107 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {

Review Comment:
   We should pass in a string as the key of the `segmentAssignmentConfigMap`, which can also be a tier name. Also suggest re-ordering the parameters a little bit as the instance partitions is optional
   ```suggestion
     public static SegmentAssignmentStrategy getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig, String assignmentType, @Nullable InstancePartitions instancePartitions) {
   ```



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyUtils.java:
##########
@@ -0,0 +1,50 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Utility class for segment assignment strategy.
+ */
+public class SegmentAssignmentStrategyUtils {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyUtils.class);
+  private SegmentAssignmentStrategyUtils() {
+  }
+
+  /**
+   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
+   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
+   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
+   * or replication changed without reassigning instances).
+   */
+  public static void checkReplication(InstancePartitions instancePartitions, int replication, String tableName) {

Review Comment:
   Suggest moving this into the `ReplicaGroupSegmentAssignmentStrategy` instead of making it a util. It does not apply to other strategies.



##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -680,6 +680,7 @@ public static class Offline {
     public static class AssignmentStrategy {
       public static final String BALANCE_NUM_SEGMENT_ASSIGNMENT_STRATEGY = "BalanceNumSegmentAssignmentStrategy";
       public static final String REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY = "ReplicaGroupSegmentAssignmentStrategy";
+      public static final String DIM_TABLE_SEGMENT_ASSIGNMENT_STRATEGY = "DimTableSegmentAssignmentStrategy";

Review Comment:
   +1 on keeping them short and concise. We might not want a strategy name for the dimension table. If we want to make one, we want to make it more general, such as `allServers`



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r980771253


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyUtils.java:
##########
@@ -0,0 +1,50 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Utility class for segment assignment strategy.
+ */
+public class SegmentAssignmentStrategyUtils {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyUtils.class);
+  private SegmentAssignmentStrategyUtils() {
+  }
+
+  /**
+   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
+   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
+   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
+   * or replication changed without reassigning instances).
+   */
+  public static void checkReplication(InstancePartitions instancePartitions, int replication, String tableName) {

Review Comment:
   done



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,107 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r968919868


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   (nit) for some reason, still seeing `_tableConfig` as a class wide variable. 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] jackjlli commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
jackjlli commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r968923302


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    _assignmentStrategyMap = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitionsMap);

Review Comment:
   The method name here seems a bit contradicting to what we are trying to do here.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r976862985


##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -680,6 +680,7 @@ public static class Offline {
     public static class AssignmentStrategy {
       public static final String BALANCE_NUM_SEGMENT_ASSIGNMENT_STRATEGY = "BalanceNumSegmentAssignmentStrategy";
       public static final String REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY = "ReplicaGroupSegmentAssignmentStrategy";
+      public static final String DIM_TABLE_SEGMENT_ASSIGNMENT_STRATEGY = "DimTableSegmentAssignmentStrategy";

Review Comment:
   also, are we covering the different cases? we can make all of the following to work. `e.g. balanced, Balanced, BaLanCed...etc`



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] codecov-commenter commented on pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
codecov-commenter commented on PR #9309:
URL: https://github.com/apache/pinot/pull/9309#issuecomment-1233441654

   # [Codecov](https://codecov.io/gh/apache/pinot/pull/9309?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) Report
   > Merging [#9309](https://codecov.io/gh/apache/pinot/pull/9309?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (145e6c9) into [master](https://codecov.io/gh/apache/pinot/commit/1c9528c1f54d90a0a62b1795c8785cf5647a676a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (1c9528c) will **decrease** coverage by `43.62%`.
   > The diff coverage is `32.19%`.
   
   ```diff
   @@              Coverage Diff              @@
   ##             master    #9309       +/-   ##
   =============================================
   - Coverage     69.73%   26.10%   -43.63%     
   + Complexity     5017       44     -4973     
   =============================================
     Files          1873     1866        -7     
     Lines         99598    99359      -239     
     Branches      15163    15141       -22     
   =============================================
   - Hits          69453    25939    -43514     
   - Misses        25217    70805    +45588     
   + Partials       4928     2615     -2313     
   ```
   
   | Flag | Coverage Δ | |
   |---|---|---|
   | integration1 | `26.10% <32.19%> (-0.04%)` | :arrow_down: |
   | integration2 | `?` | |
   | unittests1 | `?` | |
   | unittests2 | `?` | |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#carryforward-flags-in-the-pull-request-comment) to find out more.
   
   | [Impacted Files](https://codecov.io/gh/apache/pinot/pull/9309?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | Coverage Δ | |
   |---|---|---|
   | [...nt/strategy/DimTableSegmentAssignmentStrategy.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29udHJvbGxlci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29udHJvbGxlci9oZWxpeC9jb3JlL2Fzc2lnbm1lbnQvc2VnbWVudC9zdHJhdGVneS9EaW1UYWJsZVNlZ21lbnRBc3NpZ25tZW50U3RyYXRlZ3kuamF2YQ==) | `0.00% <0.00%> (ø)` | |
   | [...trategy/ReplicaGroupSegmentAssignmentStrategy.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29udHJvbGxlci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29udHJvbGxlci9oZWxpeC9jb3JlL2Fzc2lnbm1lbnQvc2VnbWVudC9zdHJhdGVneS9SZXBsaWNhR3JvdXBTZWdtZW50QXNzaWdubWVudFN0cmF0ZWd5LmphdmE=) | `0.00% <0.00%> (ø)` | |
   | [...gment/strategy/SegmentAssignmentStrategyUtils.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29udHJvbGxlci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29udHJvbGxlci9oZWxpeC9jb3JlL2Fzc2lnbm1lbnQvc2VnbWVudC9zdHJhdGVneS9TZWdtZW50QXNzaWdubWVudFN0cmF0ZWd5VXRpbHMuamF2YQ==) | `0.00% <0.00%> (ø)` | |
   | [...org/apache/pinot/spi/config/table/TableConfig.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvY29uZmlnL3RhYmxlL1RhYmxlQ29uZmlnLmphdmE=) | `0.00% <0.00%> (-64.20%)` | :arrow_down: |
   | [...nfig/table/assignment/SegmentAssignmentConfig.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvY29uZmlnL3RhYmxlL2Fzc2lnbm1lbnQvU2VnbWVudEFzc2lnbm1lbnRDb25maWcuamF2YQ==) | `0.00% <0.00%> (ø)` | |
   | [...va/org/apache/pinot/spi/utils/CommonConstants.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvdXRpbHMvQ29tbW9uQ29uc3RhbnRzLmphdmE=) | `0.00% <ø> (-27.70%)` | :arrow_down: |
   | [...he/pinot/spi/utils/builder/TableConfigBuilder.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvdXRpbHMvYnVpbGRlci9UYWJsZUNvbmZpZ0J1aWxkZXIuamF2YQ==) | `0.00% <0.00%> (-83.34%)` | :arrow_down: |
   | [...ore/assignment/segment/SegmentAssignmentUtils.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29udHJvbGxlci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29udHJvbGxlci9oZWxpeC9jb3JlL2Fzc2lnbm1lbnQvc2VnbWVudC9TZWdtZW50QXNzaWdubWVudFV0aWxzLmphdmE=) | `40.51% <14.28%> (-58.87%)` | :arrow_down: |
   | [...ent/strategy/SegmentAssignmentStrategyFactory.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29udHJvbGxlci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29udHJvbGxlci9oZWxpeC9jb3JlL2Fzc2lnbm1lbnQvc2VnbWVudC9zdHJhdGVneS9TZWdtZW50QXNzaWdubWVudFN0cmF0ZWd5RmFjdG9yeS5qYXZh) | `38.23% <38.23%> (ø)` | |
   | [...he/pinot/common/utils/config/TableConfigUtils.java](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9jb21tb24vdXRpbHMvY29uZmlnL1RhYmxlQ29uZmlnVXRpbHMuamF2YQ==) | `51.53% <40.00%> (-30.82%)` | :arrow_down: |
   | ... and [1372 more](https://codecov.io/gh/apache/pinot/pull/9309/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | |
   
   :mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation)
   


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969867675


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,135 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
+
+  private static HelixManager _helixManager;
+  private static String _tableName;
+  private static String _partitionColumn;
+  private int _replication;
+  private TableConfig _tableConfig;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _helixManager = helixManager;
+    _tableConfig = tableConfig;
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    ReplicaGroupStrategyConfig replicaGroupStrategyConfig =
+        tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();
+    _partitionColumn = replicaGroupStrategyConfig != null ? replicaGroupStrategyConfig.getPartitionColumn() : null;
+    if (_partitionColumn == null) {
+      LOGGER.info("Initialized ReplicaGroupSegmentAssignmentStrategy "
+              + "with replication: {} without partition column for table: {} ", _replication, _tableName);
+    } else {
+      LOGGER.info("Initialized ReplicaGroupSegmentAssignmentStrategy "
+              + "with replication: {} and partition column: {} for table: {}",
+              _replication, _partitionColumn, _tableName);
+    }
+  }
+
+  /**
+   * Assigns the segment for the replica-group based segment assignment strategy and returns the assigned instances.
+   */
+  @Override
+  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    int numPartitions = instancePartitions.getNumPartitions();
+    SegmentAssignmentStrategyUtils.checkReplication(instancePartitions, _replication, _tableName);
+    int partitionId;
+    if (_partitionColumn == null || numPartitions == 1) {
+      partitionId = 0;
+    } else {
+      // Uniformly spray the segment partitions over the instance partitions
+      if (_tableConfig.getTableType() == TableType.OFFLINE) {
+        partitionId =
+            SegmentAssignmentUtils
+                .getOfflineSegmentPartitionId(segmentName, _tableName, _helixManager, _partitionColumn) % numPartitions;

Review Comment:
   partitionId function is calculated different for Realtime and Offline Segment Assignments. I tried refactoring but it seems more complicated to merge



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969906651


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,87 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableNameWithType;

Review Comment:
   (nit): can be converted to local variable



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970163218


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,109 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, InstancePartitions instancePartitions) {

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,110 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    String assignmentStrategy = null;
+    Map<InstancePartitionsType, SegmentAssignmentStrategy> segmentAssignmentStrategyMap = new HashMap<>();
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    InstancePartitionsType instancePartitionsType = currentTableType == TableType.OFFLINE
+        ? InstancePartitionsType.OFFLINE
+        : InstancePartitionsType.COMPLETED;
+
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      segmentAssignmentStrategyMap.put(instancePartitionsType, segmentAssignmentStrategy);
+      return segmentAssignmentStrategyMap;
+    } else {
+      // Try to determine segment assignment strategy from table config
+      if (segmentAssignmentConfigMap != null) {
+        SegmentAssignmentConfig segmentAssignmentConfig;
+        // Use the pre defined segment assignment strategy
+        segmentAssignmentConfig = segmentAssignmentConfigMap.get(instancePartitionsType.toString());

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970399076


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   if `completedInstancePartitions  = null`, we may get `segmentAssignmentStrategy = null`. rebalanceTiers can potentially face the null pointer exception in this case?
   
   Also, is `completedInstancePartitions = null` even the valid case? If not, we can add the Precondition check.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -90,9 +89,20 @@ public List<String> assignSegment(String segmentName, Map<String, Map<String, St
     InstancePartitions instancePartitions = typeToInstancePartitions.getValue();
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
+
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;

Review Comment:
   I think that the L93-105 can be written as the following:
   
   ```
       List<String> instancesAssigned;
       if (instancePartitionsType == InstancePartitionsType.COMPLETED) {
         // Gets Segment assignment strategy for instance partitions
         SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory
             .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitions, instancePartitionsType);
         instancesAssigned = segmentAssignmentStrategy.assignSegment(segmentName,
                   currentAssignment, instancePartitions, InstancePartitionsType.COMPLETED);
       } else {
         instancesAssigned = assignConsumingSegment(SegmentName, instancePartitions);
       }
   ```
   ```



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,105 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    String assignmentStrategy = null;
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    // See https://github.com/apache/pinot/issues/9047
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      Preconditions.checkState(currentTableType == TableType.OFFLINE,
+          "DIM table Segment assignment Strategy is only applicable to OfflineTables");
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      return segmentAssignmentStrategy;
+    } else {
+      // Try to determine segment assignment strategy from table config
+      if (segmentAssignmentConfigMap != null) {
+        SegmentAssignmentConfig segmentAssignmentConfig;
+        // Use the pre defined segment assignment strategy
+        segmentAssignmentConfig = segmentAssignmentConfigMap.get(instancePartitionsType.toString().toUpperCase());
+        // Segment assignment config is only applicable to offline tables and completed segments of real time tables
+        if (segmentAssignmentConfig != null) {
+          assignmentStrategy = segmentAssignmentConfig.getAssignmentStrategy();
+        }
+      }
+    }
+
+    // Use the existing information to determine segment assignment strategy
+    SegmentAssignmentStrategy segmentAssignmentStrategy;
+    if (assignmentStrategy == null) {
+      // Calculate numReplicaGroups and numPartitions to determine segment assignment strategy
+      int numReplicaGroups = instancePartitions.getNumReplicaGroups();

Review Comment:
   I think that we will get the `null pointer error` if `instancePartitions = null`.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -53,62 +49,57 @@ protected int getReplication(TableConfig tableConfig) {
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitions, InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment, instancePartitions);
+    List<String> instancesAssigned = segmentAssignmentStrategy.assignSegment(segmentName,
+        currentAssignment, instancePartitions, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, Configuration config) {
     InstancePartitions offlineInstancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory

Review Comment:
   Do we need the null check for `offlineInstancePartitions`?
   
   



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -53,62 +49,57 @@ protected int getReplication(TableConfig tableConfig) {
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitions, InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",

Review Comment:
   Let's move this line up so that we call the precondition and then pass `instancePartitions` to the factory. Then, we will have the guarantee that `instanceParittions` is always not null.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969175475


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(

Review Comment:
   Let's remove this function. We don't need this. 
   
   we can directly call factory to get the assignment instance. No need to cache the strategy map.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    _assignmentStrategyMap = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitionsMap);

Review Comment:
   @jackjlli Can you elaborate what's contradicting?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   +1 let's delete this.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,110 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    String assignmentStrategy = null;
+    Map<InstancePartitionsType, SegmentAssignmentStrategy> segmentAssignmentStrategyMap = new HashMap<>();
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    InstancePartitionsType instancePartitionsType = currentTableType == TableType.OFFLINE
+        ? InstancePartitionsType.OFFLINE
+        : InstancePartitionsType.COMPLETED;
+
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      segmentAssignmentStrategyMap.put(instancePartitionsType, segmentAssignmentStrategy);
+      return segmentAssignmentStrategyMap;
+    } else {
+      // Try to determine segment assignment strategy from table config
+      if (segmentAssignmentConfigMap != null) {
+        SegmentAssignmentConfig segmentAssignmentConfig;
+        // Use the pre defined segment assignment strategy
+        segmentAssignmentConfig = segmentAssignmentConfigMap.get(instancePartitionsType.toString());

Review Comment:
   We should consider to standardize the key to be `toUpperCase()` so make the input config more robust.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   @GSharayu I think that @Jackie-Jiang's suggestion is simpler. 
   
   ```
   --> OfflineSegmentAssignment
   
   strategy = getSegmentAssignmentStrategy(OFFLINE)
   // deal with oflfine
     
   --> RealtimeSegmentAssignment
   
   strategy = getSegmentAssignmentStrategy(COMPLETED)
   // deal with completed
   
   ...
   strategy = getSegmentAssignmentStrategy(CONSUMING)
   // deal with consuming
   
   ```



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] jackjlli commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
jackjlli commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r968946171


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,110 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    String assignmentStrategy = null;
+    Map<InstancePartitionsType, SegmentAssignmentStrategy> segmentAssignmentStrategyMap = new HashMap<>();
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    InstancePartitionsType instancePartitionsType = currentTableType == TableType.OFFLINE
+        ? InstancePartitionsType.OFFLINE
+        : InstancePartitionsType.COMPLETED;
+
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      segmentAssignmentStrategyMap.put(instancePartitionsType, segmentAssignmentStrategy);

Review Comment:
   Should we add a precondition check here to make sure whether the instancePartitionType is `OFFLINE` not `COMPLETED`?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -388,4 +395,89 @@ public Map<String, Map<String, String>> getNonTierSegmentAssignment() {
       return _nonTierSegmentAssignment;
     }
   }
+
+  /**
+   * Returns a partition id for offline table
+   */
+  public static int getOfflineSegmentPartitionId(String segmentName, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    SegmentZKMetadata segmentZKMetadata =
+        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName, segmentName);
+    Preconditions.checkState(segmentZKMetadata != null,
+        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, offlineTableName);
+    return getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+  }
+
+  private static int getPartitionId(SegmentZKMetadata segmentZKMetadata,
+      String offlineTableName, @Nullable String partitionColumn) {

Review Comment:
   Does it have to be offline table?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -388,4 +395,89 @@ public Map<String, Map<String, String>> getNonTierSegmentAssignment() {
       return _nonTierSegmentAssignment;
     }
   }
+
+  /**
+   * Returns a partition id for offline table
+   */
+  public static int getOfflineSegmentPartitionId(String segmentName, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    SegmentZKMetadata segmentZKMetadata =
+        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName, segmentName);
+    Preconditions.checkState(segmentZKMetadata != null,
+        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, offlineTableName);
+    return getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+  }
+
+  private static int getPartitionId(SegmentZKMetadata segmentZKMetadata,
+      String offlineTableName, @Nullable String partitionColumn) {
+    String segmentName = segmentZKMetadata.getSegmentName();
+    ColumnPartitionMetadata partitionMetadata =
+        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(partitionColumn);
+    Preconditions.checkState(partitionMetadata != null,
+        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
+        segmentName, offlineTableName, partitionColumn);
+    Set<Integer> partitions = partitionMetadata.getPartitions();
+    Preconditions.checkState(partitions.size() == 1,
+        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
+        offlineTableName, partitionColumn);
+    return partitions.iterator().next();
+  }
+
+  /**
+   * Returns map of instance partition id to segments for offline tables
+   */
+  public static Map<Integer, List<String>> getOfflineInstancePartitionIdToSegmentsMap(Set<String> segments,
+      int numInstancePartitions, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    // Fetch partition id from segment ZK metadata
+    List<SegmentZKMetadata> segmentsZKMetadata =
+        ZKMetadataProvider.getSegmentsZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName);
+
+    Map<Integer, List<String>> instancePartitionIdToSegmentsMap = new HashMap<>();
+    Set<String> segmentsWithoutZKMetadata = new HashSet<>(segments);
+    for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) {
+      String segmentName = segmentZKMetadata.getSegmentName();
+      if (segmentsWithoutZKMetadata.remove(segmentName)) {
+        int partitionId = getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+        int instancePartitionId = partitionId % numInstancePartitions;
+        instancePartitionIdToSegmentsMap.computeIfAbsent(instancePartitionId, k -> new ArrayList<>()).add(segmentName);
+      }
+    }
+    Preconditions.checkState(segmentsWithoutZKMetadata.isEmpty(), "Failed to find ZK metadata for segments: %s",
+        segmentsWithoutZKMetadata);
+
+    return instancePartitionIdToSegmentsMap;
+  }
+
+  /**
+   * Returns a partition id for realtime table
+   */
+  public static int getRealtimeSegmentPartitionId(String segmentName, String realtimeTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    Integer segmentPartitionId =
+        SegmentUtils.getRealtimeSegmentPartitionId(segmentName, realtimeTableName, helixManager, partitionColumn);
+    if (segmentPartitionId == null) {
+      // This case is for the uploaded segments for which there's no partition information.
+      // A random, but consistent, partition id is calculated based on the hash code of the segment name.
+      // Note that '% 10K' is used to prevent having partition ids with large value which will be problematic later in
+      // instance assignment formula.
+      segmentPartitionId = Math.abs(segmentName.hashCode() % 10_000);

Review Comment:
   Shouldn't we reuse one of the partition functions like `HashCodePartitionFunction` instead of maintaining one more logic to generate partition Id here?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>

Review Comment:
   The javadoc here says that it puts into a map of `Map<InstancePartitionsType, SegmentAssignmentStrategy>`, but the parameter is a map of `Map<InstancePartitionsType, InstancePartitions>`, sth that needs to be updated.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,135 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
+
+  private static HelixManager _helixManager;
+  private static String _tableName;
+  private static String _partitionColumn;
+  private int _replication;
+  private TableConfig _tableConfig;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _helixManager = helixManager;
+    _tableConfig = tableConfig;
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    ReplicaGroupStrategyConfig replicaGroupStrategyConfig =
+        tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();

Review Comment:
   nit: define one local variable to store the validation config since it's been fetched 2 times within a method.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();

Review Comment:
   do we want null check for validation config here?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized BalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableName, _replication);
+  }
+
+  @Override
+  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    validateSegmentAssignmentStrategy(instancePartitions);
+    return SegmentAssignmentUtils
+        .assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions, _replication);
+  }
+
+  @Override
+  public Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    validateSegmentAssignmentStrategy(instancePartitions);
+    Map<String, Map<String, String>> newAssignment;
+    List<String> instances =
+        SegmentAssignmentUtils.getInstancesForNonReplicaGroupBasedAssignment(instancePartitions, _replication);
+    newAssignment = SegmentAssignmentUtils
+        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, instances, _replication);
+    return newAssignment;
+  }
+
+  private void validateSegmentAssignmentStrategy(InstancePartitions instancePartitions) {
+    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
+    int numPartitions = instancePartitions.getNumPartitions();
+
+    Preconditions.checkState(numReplicaGroups == 1, "Replica groups should be 1");

Review Comment:
   Could you add some javadoc to the method on why these two configs must be 1? Or modify the error message here to explain more details like `Replica groups should be 1 in order to use BalanceNumSegmentAssignmentStrategy`?



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;

Review Comment:
   nit: rename it to `_tableNameWithType` here since the value is fetched from a table config.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -125,8 +133,13 @@ private List<String> assignConsumingSegment(String segmentName, InstancePartitio
       return instancesAssigned;
     } else {
       // Replica-group based assignment
-
-      checkReplication(instancePartitions);
+      // TODO: Refactor check replication this for segment assignment strategy in follow up PR

Review Comment:
   nit: link the OSS issue here. Same for other places with `TODO`.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,135 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
+
+  private static HelixManager _helixManager;
+  private static String _tableName;
+  private static String _partitionColumn;
+  private int _replication;
+  private TableConfig _tableConfig;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _helixManager = helixManager;
+    _tableConfig = tableConfig;
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    ReplicaGroupStrategyConfig replicaGroupStrategyConfig =
+        tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();
+    _partitionColumn = replicaGroupStrategyConfig != null ? replicaGroupStrategyConfig.getPartitionColumn() : null;
+    if (_partitionColumn == null) {
+      LOGGER.info("Initialized ReplicaGroupSegmentAssignmentStrategy "
+              + "with replication: {} without partition column for table: {} ", _replication, _tableName);
+    } else {
+      LOGGER.info("Initialized ReplicaGroupSegmentAssignmentStrategy "
+              + "with replication: {} and partition column: {} for table: {}",
+              _replication, _partitionColumn, _tableName);
+    }
+  }
+
+  /**
+   * Assigns the segment for the replica-group based segment assignment strategy and returns the assigned instances.
+   */
+  @Override
+  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    int numPartitions = instancePartitions.getNumPartitions();
+    SegmentAssignmentStrategyUtils.checkReplication(instancePartitions, _replication, _tableName);
+    int partitionId;
+    if (_partitionColumn == null || numPartitions == 1) {
+      partitionId = 0;
+    } else {
+      // Uniformly spray the segment partitions over the instance partitions
+      if (_tableConfig.getTableType() == TableType.OFFLINE) {
+        partitionId =
+            SegmentAssignmentUtils
+                .getOfflineSegmentPartitionId(segmentName, _tableName, _helixManager, _partitionColumn) % numPartitions;

Review Comment:
   Same here, we have a bunch of partition functions to generate the partition Ids already. We should consider reusing them if possible.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    _assignmentStrategyMap = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitionsMap);

Review Comment:
   The method name here seems a bit contradict to what we are trying to do here.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967486511


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -183,8 +176,10 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
     // Rebalance tiers first
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.COMPLETED);

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignment.java:
##########
@@ -31,20 +31,17 @@
 
 /**
  * Interface for segment assignment and table rebalance.
- * <p>
- * TODO: Add SegmentAssignmentStrategy interface and support custom segment assignment strategy (e.g. cost based segment
- *       assignment). SegmentAssignmentStrategy should not be coupled with SegmentAssignment, and SegmentAssignment
- *       should be able to choose the segment assignment strategy based on the configuration.
  */
 public interface SegmentAssignment {
 
   /**
    * Initializes the segment assignment.
-   *
    * @param helixManager Helix manager
    * @param tableConfig Table config
+   * @param instancePartitionsMap Map from type (OFFLINE|CONSUMING|COMPLETED) to instance partitions
    */
-  void init(HelixManager helixManager, TableConfig tableConfig);
+  void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap);

Review Comment:
   updated this



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969860379


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -388,4 +395,89 @@ public Map<String, Map<String, String>> getNonTierSegmentAssignment() {
       return _nonTierSegmentAssignment;
     }
   }
+
+  /**
+   * Returns a partition id for offline table
+   */
+  public static int getOfflineSegmentPartitionId(String segmentName, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    SegmentZKMetadata segmentZKMetadata =
+        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName, segmentName);
+    Preconditions.checkState(segmentZKMetadata != null,
+        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, offlineTableName);
+    return getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+  }
+
+  private static int getPartitionId(SegmentZKMetadata segmentZKMetadata,
+      String offlineTableName, @Nullable String partitionColumn) {
+    String segmentName = segmentZKMetadata.getSegmentName();
+    ColumnPartitionMetadata partitionMetadata =
+        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(partitionColumn);
+    Preconditions.checkState(partitionMetadata != null,
+        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
+        segmentName, offlineTableName, partitionColumn);
+    Set<Integer> partitions = partitionMetadata.getPartitions();
+    Preconditions.checkState(partitions.size() == 1,
+        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
+        offlineTableName, partitionColumn);
+    return partitions.iterator().next();
+  }
+
+  /**
+   * Returns map of instance partition id to segments for offline tables
+   */
+  public static Map<Integer, List<String>> getOfflineInstancePartitionIdToSegmentsMap(Set<String> segments,
+      int numInstancePartitions, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    // Fetch partition id from segment ZK metadata
+    List<SegmentZKMetadata> segmentsZKMetadata =
+        ZKMetadataProvider.getSegmentsZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName);
+
+    Map<Integer, List<String>> instancePartitionIdToSegmentsMap = new HashMap<>();
+    Set<String> segmentsWithoutZKMetadata = new HashSet<>(segments);
+    for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) {
+      String segmentName = segmentZKMetadata.getSegmentName();
+      if (segmentsWithoutZKMetadata.remove(segmentName)) {
+        int partitionId = getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+        int instancePartitionId = partitionId % numInstancePartitions;
+        instancePartitionIdToSegmentsMap.computeIfAbsent(instancePartitionId, k -> new ArrayList<>()).add(segmentName);
+      }
+    }
+    Preconditions.checkState(segmentsWithoutZKMetadata.isEmpty(), "Failed to find ZK metadata for segments: %s",
+        segmentsWithoutZKMetadata);
+
+    return instancePartitionIdToSegmentsMap;
+  }
+
+  /**
+   * Returns a partition id for realtime table
+   */
+  public static int getRealtimeSegmentPartitionId(String segmentName, String realtimeTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    Integer segmentPartitionId =
+        SegmentUtils.getRealtimeSegmentPartitionId(segmentName, realtimeTableName, helixManager, partitionColumn);
+    if (segmentPartitionId == null) {
+      // This case is for the uploaded segments for which there's no partition information.
+      // A random, but consistent, partition id is calculated based on the hash code of the segment name.
+      // Note that '% 10K' is used to prevent having partition ids with large value which will be problematic later in
+      // instance assignment formula.
+      segmentPartitionId = Math.abs(segmentName.hashCode() % 10_000);

Review Comment:
   I will scope out comment from this PR as the change was discussed in https://github.com/apache/pinot/pull/8584



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969907045


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +67,14 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;

Review Comment:
   Looks like not needed anymore



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r976047082


##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -680,6 +680,7 @@ public static class Offline {
     public static class AssignmentStrategy {
       public static final String BALANCE_NUM_SEGMENT_ASSIGNMENT_STRATEGY = "BalanceNumSegmentAssignmentStrategy";
       public static final String REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY = "ReplicaGroupSegmentAssignmentStrategy";
+      public static final String DIM_TABLE_SEGMENT_ASSIGNMENT_STRATEGY = "DimTableSegmentAssignmentStrategy";

Review Comment:
   Our new config doesn't necessarily need to use the same pattern. Alternative way is to have sth like the following:
   
   ```
   {
     "segmentAssignmentConfigMap": {
       "OFFLINE": {
         "segmentAssignmentStrategy": "balanced/replicaGroup/dimTable"
       }
     }
     ...
   }
   
   ```



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r980774073


##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -680,6 +680,7 @@ public static class Offline {
     public static class AssignmentStrategy {
       public static final String BALANCE_NUM_SEGMENT_ASSIGNMENT_STRATEGY = "BalanceNumSegmentAssignmentStrategy";
       public static final String REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY = "ReplicaGroupSegmentAssignmentStrategy";
+      public static final String DIM_TABLE_SEGMENT_ASSIGNMENT_STRATEGY = "DimTableSegmentAssignmentStrategy";

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,87 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on PR #9309:
URL: https://github.com/apache/pinot/pull/9309#issuecomment-1241475640

   ```2022-09-08T23:47:33.2813064Z [ERROR] org.apache.pinot.controller.helix.core.rebalance.TableRebalancerClusterStatelessTest.testRebalanceWithTiers  Time elapsed: 5.094 s  <<< FAILURE!
   2022-09-08T23:47:33.2814091Z org.apache.pinot.controller.api.exception.TableAlreadyExistsException: Table config for testTable_OFFLINE already exists. If this is unexpected, try deleting the table to remove all metadata associated with it.
   2022-09-08T23:47:33.2815139Z 	at org.apache.pinot.controller.helix.core.PinotHelixResourceManager.addTable(PinotHelixResourceManager.java:1412)
   2022-09-08T23:47:33.2816081Z 	at org.apache.pinot.controller.helix.core.rebalance.TableRebalancerClusterStatelessTest.testRebalanceWithTiers(TableRebalancerClusterStatelessTest.java:343)
   2022-09-08T23:47:33.2816919Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
   2022-09-08T23:47:33.2817485Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
   2022-09-08T23:47:33.2818120Z 	at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
   2022-09-08T23:47:33.2818640Z 	at java.base/java.lang.reflect.Method.invoke(Method.java:566)
   2022-09-08T23:47:33.2819129Z 	at org.testng.internal.MethodInvocationHelper.invokeMethod(MethodInvocationHelper.java:108)
   2022-09-08T23:47:33.2819624Z 	at org.testng.internal.Invoker.invokeMethod(Invoker.java:661)
   2022-09-08T23:47:33.2820029Z 	at org.testng.internal.Invoker.invokeTestMethod(Invoker.java:869)
   2022-09-08T23:47:33.2820469Z 	at org.testng.internal.Invoker.invokeTestMethods(Invoker.java:1193)
   2022-09-08T23:47:33.2820970Z 	at org.testng.internal.TestMethodWorker.invokeTestMethods(TestMethodWorker.java:126)
   2022-09-08T23:47:33.2835876Z 	at org.testng.internal.TestMethodWorker.run(TestMethodWorker.java:109)
   2022-09-08T23:47:33.2836443Z 	at org.testng.TestRunner.privateRun(TestRunner.java:744)
   2022-09-08T23:47:33.2836810Z 	at org.testng.TestRunner.run(TestRunner.java:602)
   2022-09-08T23:47:33.2837169Z 	at org.testng.SuiteRunner.runTest(SuiteRunner.java:380)
   2022-09-08T23:47:33.2837552Z 	at org.testng.SuiteRunner.runSequentially(SuiteRunner.java:375)
   2022-09-08T23:47:33.2837963Z 	at org.testng.SuiteRunner.privateRun(SuiteRunner.java:340)
   2022-09-08T23:47:33.2838328Z 	at org.testng.SuiteRunner.run(SuiteRunner.java:289)
   2022-09-08T23:47:33.2838718Z 	at org.testng.SuiteRunnerWorker.runSuite(SuiteRunnerWorker.java:52)
   2022-09-08T23:47:33.2839140Z 	at org.testng.SuiteRunnerWorker.run(SuiteRunnerWorker.java:86)
   2022-09-08T23:47:33.2839547Z 	at org.testng.TestNG.runSuitesSequentially(TestNG.java:1301)
   2022-09-08T23:47:33.2839920Z 	at org.testng.TestNG.runSuitesLocally(TestNG.java:1226)
   2022-09-08T23:47:33.2840271Z 	at org.testng.TestNG.runSuites(TestNG.java:1144)
   2022-09-08T23:47:33.2840588Z 	at org.testng.TestNG.run(TestNG.java:1115)
   2022-09-08T23:47:33.2841002Z 	at org.apache.maven.surefire.testng.TestNGExecutor.run(TestNGExecutor.java:284)
   2022-09-08T23:47:33.2841558Z 	at org.apache.maven.surefire.testng.TestNGXmlTestSuite.execute(TestNGXmlTestSuite.java:75)
   2022-09-08T23:47:33.2842129Z 	at org.apache.maven.surefire.testng.TestNGProvider.invoke(TestNGProvider.java:119)
   2022-09-08T23:47:33.2842887Z 	at org.apache.maven.surefire.booter.ForkedBooter.runSuitesInProcess(ForkedBooter.java:428)
   2022-09-08T23:47:33.2843451Z 	at org.apache.maven.surefire.booter.ForkedBooter.execute(ForkedBooter.java:162)
   2022-09-08T23:47:33.2843953Z 	at org.apache.maven.surefire.booter.ForkedBooter.run(ForkedBooter.java:562)
   2022-09-08T23:47:33.2844449Z 	at org.apache.maven.surefire.booter.ForkedBooter.main(ForkedBooter.java:548)
   2022-09-08T23:47:33.2844718Z 
   2022-09-08T23:47:33.2845206Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testUpdateBrokerResource  Time elapsed: 0.216 s
   2022-09-08T23:47:33.2846308Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateDimTableTenantConfig  Time elapsed: 0.072 s
   2022-09-08T23:47:33.2847330Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateTenantConfig  Time elapsed: 0.052 s
   2022-09-08T23:47:33.2848188Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testBrokerTenant  Time elapsed: 0.953 s
   2022-09-08T23:47:33.2848940Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testEmptyServerTenant  Time elapsed: 0.02 s
   2022-09-08T23:47:33.2849714Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testServerTenant  Time elapsed: 0.388 s
   2022-09-08T23:47:33.2850500Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testDualModeController  Time elapsed: 9.822 s
   2022-09-08T23:47:33.2851344Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testHelixOnlyController  Time elapsed: 1.718 s
   2022-09-08T23:47:33.2852184Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testPinotOnlyController  Time elapsed: 5.248 s
   2022-09-08T23:47:33.2853009Z [ERROR] org.apache.pinot.controller.api.PinotBrokerRestletResourceStatelessTest.testGetBrokers  Time elapsed: 14.044 s
   2022-09-08T23:47:33.2853853Z [ERROR] org.apache.pinot.controller.api.PinotIngestionRestletResourceStatelessTest.testIngestEndpoint  Time elapsed: 0.492 s
   2022-09-08T23:47:33.2855076Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testDefaultPinotTaskManagerNoScheduler  Time elapsed: 3.905 s
   2022-09-08T23:47:33.2856187Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithRestart  Time elapsed: 8.857 s
   2022-09-08T23:47:33.2857313Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithUpdate  Time elapsed: 2.65 s
   2022-09-08T23:47:33.2858407Z [ERROR] org.apache.pinot.controller.validation.ValidationManagerStatelessTest.testRebuildBrokerResourceWhenBrokerAdded  Time elapsed: 0.04 s
   2022-09-08T23:47:33.3510290Z [INFO] 
   2022-09-08T23:47:33.3511030Z [INFO] Results:
   2022-09-08T23:47:33.3512144Z [INFO] 
   2022-09-08T23:47:33.3512504Z [ERROR] Failures: 
   2022-09-08T23:47:33.3513121Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalance:210 expected [DONE] but found [FAILED]
   2022-09-08T23:47:33.3534306Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalanceWithTiers:343 » TableAlreadyExists
   ```
   
   It looks that rebalancing is failing.


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969850739


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   Not sure if I still understand this comment, we are using `_tableConfig ` in setSegmentAssignmentStrategyMap which is called from both RealtimeSegmentAssignment and OfflineSegmentAssignment.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r964211022


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -44,6 +40,13 @@
  */
 public class OfflineSegmentAssignment extends BaseSegmentAssignment {

Review Comment:
   (nit), while it may not make sense for `RealtimeSegmentAssignment`, for `OfflineSegmentAssignment`, we can consider making `InstancePartitionsType.OFFLINE` an instance variable as it is referenced 9 times in this class.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r966626344


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigUtils.java:
##########
@@ -165,9 +166,18 @@ public static TableConfig fromZNRecord(ZNRecord znRecord)
           new TypeReference<Map<InstancePartitionsType, String>>() { });
     }
 
+    Map<InstancePartitionsType, SegmentAssignmentConfig> segmentAssignmentConfigMap = null;
+    String segmentAssignmentConfigMapString = simpleFields.get(TableConfig.SEGMENT_ASSIGNMENT_CONFIG_MAP_KEY);
+    if (segmentAssignmentConfigMapString != null) {
+      segmentAssignmentConfigMap = JsonUtils.stringToObject(segmentAssignmentConfigMapString,
+          new TypeReference<>() {

Review Comment:
   Compatibility test is complaining the following:
   ```
   [INFO] -------------------------------------------------------------
   Error:  COMPILATION ERROR : 
   [INFO] -------------------------------------------------------------
   Error:  /home/runner/work/pinot/pinot/pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigUtils.java:[173,27] error: cannot infer type arguments for TypeReference<T>
     reason: cannot use '<>' with anonymous inner classes
     where T is a type-variable:
       T extends Object declared in class TypeReference
   ```
   
   It looks that we need to provide the Map types correctly. 
   ` new TypeReference<>() ->  new TypeReference<Map<InstanceParttitionsType,String>>()`



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu closed pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu closed pull request #9309: Introduce Segment AssignmentStrategy Interface 
URL: https://github.com/apache/pinot/pull/9309


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969962134


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,109 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, InstancePartitions instancePartitions) {

Review Comment:
   Don't we need to pass the `instanceType` as well?



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r971256430


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -53,62 +49,57 @@ protected int getReplication(TableConfig tableConfig) {
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitions, InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment, instancePartitions);
+    List<String> instancesAssigned = segmentAssignmentStrategy.assignSegment(segmentName,
+        currentAssignment, instancePartitions, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, Configuration config) {
     InstancePartitions offlineInstancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    // Gets Segment assignment strategy for instance partitions
+    SegmentAssignmentStrategy segmentAssignmentStrategy = SegmentAssignmentStrategyFactory

Review Comment:
   we already have that, added the same preconditions check before



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r971262433


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   There are already set of preconditions to check not null before reassigning segments and tiers, so we are not expected to see completedInstancePartitions = null, I have added precondition check in segment assignment strategy factory 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r961196682


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -75,6 +74,14 @@
  */
 public class RealtimeSegmentAssignment extends BaseSegmentAssignment {
 
+  @Override
+  protected void setSegmentAssignmentStrategyMap(

Review Comment:
   can consider moving this implementation up to be implemented in `BaseSegmentAssignment` since `OfflineSegmentAssignment` also shares same implementation and looks like `SegmentAssignmentStrategyFactory.getSegmentAssignmentStrategy` can handle the specifics. 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu closed pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu closed pull request #9309: Introduce Segment AssignmentStrategy Interface 
URL: https://github.com/apache/pinot/pull/9309


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967315539


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignment.java:
##########
@@ -31,20 +31,17 @@
 
 /**
  * Interface for segment assignment and table rebalance.
- * <p>
- * TODO: Add SegmentAssignmentStrategy interface and support custom segment assignment strategy (e.g. cost based segment
- *       assignment). SegmentAssignmentStrategy should not be coupled with SegmentAssignment, and SegmentAssignment
- *       should be able to choose the segment assignment strategy based on the configuration.
  */
 public interface SegmentAssignment {
 
   /**
    * Initializes the segment assignment.
-   *
    * @param helixManager Helix manager
    * @param tableConfig Table config
+   * @param instancePartitionsMap Map from type (OFFLINE|CONSUMING|COMPLETED) to instance partitions
    */
-  void init(HelixManager helixManager, TableConfig tableConfig);
+  void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap);

Review Comment:
   > ```
   > 2022-09-08T23:47:33.2814091Z org.apache.pinot.controller.api.exception.TableAlreadyExistsException: Table config for testTable_OFFLINE already exists. If this is unexpected, try deleting the table to remove all metadata associated with it.
   > 2022-09-08T23:47:33.2815139Z 	at org.apache.pinot.controller.helix.core.PinotHelixResourceManager.addTable(PinotHelixResourceManager.java:1412)
   > 2022-09-08T23:47:33.2816081Z 	at org.apache.pinot.controller.helix.core.rebalance.TableRebalancerClusterStatelessTest.testRebalanceWithTiers(TableRebalancerClusterStatelessTest.java:343)
   > 2022-09-08T23:47:33.2816919Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
   > 2022-09-08T23:47:33.2817485Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
   > 2022-09-08T23:47:33.2818120Z 	at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
   > 2022-09-08T23:47:33.2818640Z 	at java.base/java.lang.reflect.Method.invoke(Method.java:566)
   > 2022-09-08T23:47:33.2819129Z 	at org.testng.internal.MethodInvocationHelper.invokeMethod(MethodInvocationHelper.java:108)
   > 2022-09-08T23:47:33.2819624Z 	at org.testng.internal.Invoker.invokeMethod(Invoker.java:661)
   > 2022-09-08T23:47:33.2820029Z 	at org.testng.internal.Invoker.invokeTestMethod(Invoker.java:869)
   > 2022-09-08T23:47:33.2820469Z 	at org.testng.internal.Invoker.invokeTestMethods(Invoker.java:1193)
   > 2022-09-08T23:47:33.2820970Z 	at org.testng.internal.TestMethodWorker.invokeTestMethods(TestMethodWorker.java:126)
   > 2022-09-08T23:47:33.2835876Z 	at org.testng.internal.TestMethodWorker.run(TestMethodWorker.java:109)
   > 2022-09-08T23:47:33.2836443Z 	at org.testng.TestRunner.privateRun(TestRunner.java:744)
   > 2022-09-08T23:47:33.2836810Z 	at org.testng.TestRunner.run(TestRunner.java:602)
   > 2022-09-08T23:47:33.2837169Z 	at org.testng.SuiteRunner.runTest(SuiteRunner.java:380)
   > 2022-09-08T23:47:33.2837552Z 	at org.testng.SuiteRunner.runSequentially(SuiteRunner.java:375)
   > 2022-09-08T23:47:33.2837963Z 	at org.testng.SuiteRunner.privateRun(SuiteRunner.java:340)
   > 2022-09-08T23:47:33.2838328Z 	at org.testng.SuiteRunner.run(SuiteRunner.java:289)
   > 2022-09-08T23:47:33.2838718Z 	at org.testng.SuiteRunnerWorker.runSuite(SuiteRunnerWorker.java:52)
   > 2022-09-08T23:47:33.2839140Z 	at org.testng.SuiteRunnerWorker.run(SuiteRunnerWorker.java:86)
   > 2022-09-08T23:47:33.2839547Z 	at org.testng.TestNG.runSuitesSequentially(TestNG.java:1301)
   > 2022-09-08T23:47:33.2839920Z 	at org.testng.TestNG.runSuitesLocally(TestNG.java:1226)
   > 2022-09-08T23:47:33.2840271Z 	at org.testng.TestNG.runSuites(TestNG.java:1144)
   > 2022-09-08T23:47:33.2840588Z 	at org.testng.TestNG.run(TestNG.java:1115)
   > 2022-09-08T23:47:33.2841002Z 	at org.apache.maven.surefire.testng.TestNGExecutor.run(TestNGExecutor.java:284)
   > 2022-09-08T23:47:33.2841558Z 	at org.apache.maven.surefire.testng.TestNGXmlTestSuite.execute(TestNGXmlTestSuite.java:75)
   > 2022-09-08T23:47:33.2842129Z 	at org.apache.maven.surefire.testng.TestNGProvider.invoke(TestNGProvider.java:119)
   > 2022-09-08T23:47:33.2842887Z 	at org.apache.maven.surefire.booter.ForkedBooter.runSuitesInProcess(ForkedBooter.java:428)
   > 2022-09-08T23:47:33.2843451Z 	at org.apache.maven.surefire.booter.ForkedBooter.execute(ForkedBooter.java:162)
   > 2022-09-08T23:47:33.2843953Z 	at org.apache.maven.surefire.booter.ForkedBooter.run(ForkedBooter.java:562)
   > 2022-09-08T23:47:33.2844449Z 	at org.apache.maven.surefire.booter.ForkedBooter.main(ForkedBooter.java:548)
   > 2022-09-08T23:47:33.2844718Z 
   > 2022-09-08T23:47:33.2845206Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testUpdateBrokerResource  Time elapsed: 0.216 s
   > 2022-09-08T23:47:33.2846308Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateDimTableTenantConfig  Time elapsed: 0.072 s
   > 2022-09-08T23:47:33.2847330Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateTenantConfig  Time elapsed: 0.052 s
   > 2022-09-08T23:47:33.2848188Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testBrokerTenant  Time elapsed: 0.953 s
   > 2022-09-08T23:47:33.2848940Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testEmptyServerTenant  Time elapsed: 0.02 s
   > 2022-09-08T23:47:33.2849714Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testServerTenant  Time elapsed: 0.388 s
   > 2022-09-08T23:47:33.2850500Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testDualModeController  Time elapsed: 9.822 s
   > 2022-09-08T23:47:33.2851344Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testHelixOnlyController  Time elapsed: 1.718 s
   > 2022-09-08T23:47:33.2852184Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testPinotOnlyController  Time elapsed: 5.248 s
   > 2022-09-08T23:47:33.2853009Z [ERROR] org.apache.pinot.controller.api.PinotBrokerRestletResourceStatelessTest.testGetBrokers  Time elapsed: 14.044 s
   > 2022-09-08T23:47:33.2853853Z [ERROR] org.apache.pinot.controller.api.PinotIngestionRestletResourceStatelessTest.testIngestEndpoint  Time elapsed: 0.492 s
   > 2022-09-08T23:47:33.2855076Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testDefaultPinotTaskManagerNoScheduler  Time elapsed: 3.905 s
   > 2022-09-08T23:47:33.2856187Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithRestart  Time elapsed: 8.857 s
   > 2022-09-08T23:47:33.2857313Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithUpdate  Time elapsed: 2.65 s
   > 2022-09-08T23:47:33.2858407Z [ERROR] org.apache.pinot.controller.validation.ValidationManagerStatelessTest.testRebuildBrokerResourceWhenBrokerAdded  Time elapsed: 0.04 s
   > 2022-09-08T23:47:33.3510290Z [INFO] 
   > 2022-09-08T23:47:33.3511030Z [INFO] Results:
   > 2022-09-08T23:47:33.3512144Z [INFO] 
   > 2022-09-08T23:47:33.3512504Z [ERROR] Failures: 
   > 2022-09-08T23:47:33.3513121Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalance:210 expected [DONE] but found [FAILED]
   > 2022-09-08T23:47:33.3534306Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalanceWithTiers:343 » TableAlreadyExists
   > ```
   > 
   > It looks that rebalancing is failing.
   
   Fixed this as well.I think this was not showing up because of other flaky test.  In the existing code, we set BalanceNumSegmentAssignmentStrategy as default and update table config replication which did not match the assignment strategy which caused this test to fail. Handled change in strategyAssignmentFactory



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967313797


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -194,60 +171,16 @@ protected Map<String, Map<String, String>> reassignSegments(String instanceParti
       // When bootstrap is enabled, start with an empty assignment and reassign all segments
       newAssignment = new TreeMap<>();
       for (String segment : currentAssignment.keySet()) {
-        List<String> assignedInstances = assignSegment(segment, newAssignment, instancePartitions);
+        List<String> assignedInstances = assignSegment(segment, newAssignment, instancePartitions,

Review Comment:
   Updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967486246


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r965283986


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment, instancePartitions);
+    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment,
+            instancePartitions, segmentAssignmentStrategy, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, Configuration config) {
     InstancePartitions offlineInstancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.reassignSegments(currentAssignment,
+          offlineInstancePartitions, InstancePartitionsType.OFFLINE);

Review Comment:
   For Dim tables, we would have reassignSegments bypassing bootstrap, so should be fine



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r965280745


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,67 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+
+
+/**
+ * Interface for segment assignment strategies
+ */
+public interface SegmentAssignmentStrategy {
+
+  /**
+   * Initializes the segment assignment strategy.
+   *
+   * @param helixManager Helix manager
+   * @param tableConfig Table config
+   */
+  void init(HelixManager helixManager, TableConfig tableConfig);
+
+  /**
+   * Assigns segment to instances. The assignment strategy will be configured in
+   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
+   * assignment strategy, this function will be called to assign a new segment
+   *
+   * @param segmentName Name of the segment to be assigned
+   * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
+   * @param instancePartitions Instance partitions
+   * @return List of instances to assign the segment to
+   */
+  List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType);
+
+  /**
+   * Re-assigns segment to instances. The assignment strategy will be configured in
+   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
+   * assignment strategy, this function will be called to re-assign a segment
+   * when the InstancePartitions has been changed.
+   *
+   * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
+   * @param instancePartitions Instance partitions
+   * @return

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] klsince commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
klsince commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r1018496195


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   hi @GSharayu , I'm testing locally with RealtimeQuickStart but customized a little bit, and I came into NPE issue like below. The tableConfig has some tierConfigs, and SegmentRelocator is enabled to run periodically to move segments around based on data age as set in tierConfigs. As suggested above, I feel it'd be safer to init segmentAssignmentStrategy with a default policy. 
   
   ```
   2022/11/09 14:43:09.474 INFO [RealtimeSegmentAssignment] [async-task-thread-1] Rebalancing tier: myTier for table: airlineStats_REALTIME with bootstrap: false, instance partitions: {"instancePartitionsName":"airlineStats
   _myTier","partitionToInstancesMap":{"0_0":["Server_192.168.0.106_7000"]}}
   2022/11/09 14:43:09.475 WARN [TableRebalancer] [async-task-thread-1] Caught exception while calculating target assignment for table: airlineStats_REALTIME, aborting the rebalance
   java.lang.NullPointerException: null
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.reassignSegments(BaseSegmentAssignment.java:158) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.rebalanceTiers(BaseSegmentAssignment.java:129) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.RealtimeSegmentAssignment.rebalanceTable(RealtimeSegmentAssignment.java:201) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.rebalance.TableRebalancer.rebalance(TableRebalancer.java:219) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.relocation.SegmentRelocator.lambda$processTable$0(SegmentRelocator.java:137) ~[classes/:?]
           at java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:515) [?:?]
           at java.util.concurrent.FutureTask.run(FutureTask.java:264) [?:?]
           at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128) [?:?]
           at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628) [?:?]
           at java.lang.Thread.run(Thread.java:829) [?:?]
   ```
   
   cc @npawar for awareness ^



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r964221907


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);

Review Comment:
   ```suggestion
       InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
       Preconditions.checkState(
           !(segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) && instancePartitions != null,
           "Failed to find OFFLINE instance partitions for table: %s", _tableNameWithType);
   ```
   One way of rewriting to not bypass the logging below and avoid calling `segmentAssignmentStrategy.reassignSegments` directly.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on PR #9309:
URL: https://github.com/apache/pinot/pull/9309#issuecomment-1234929411

   `SegmentAssignment`'s TODO can probably be taken out or modified as this PR addresses the `SegmentAssignmentStrategy` part. 


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r964225345


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment, instancePartitions);
+    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment,
+            instancePartitions, segmentAssignmentStrategy, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, Configuration config) {
     InstancePartitions offlineInstancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.reassignSegments(currentAssignment,
+          offlineInstancePartitions, InstancePartitionsType.OFFLINE);

Review Comment:
   We should call `reassignSegments(...)` instead of  `segmentAssignmentStrategy.reassignSegments(...)` as this will introduce different behavior (e.g. bypass bootstrap). Feel free to consider the suggestion above ^. 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r965283050


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -185,7 +157,8 @@ protected Pair<List<Map<String, Map<String, String>>>, Map<String, Map<String, S
    * Rebalances segments in the current assignment using the instancePartitions and returns new assignment
    */
   protected Map<String, Map<String, String>> reassignSegments(String instancePartitionType,
-      Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, boolean bootstrap) {
+      Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, boolean bootstrap,
+      SegmentAssignmentStrategy segmentAssignmentStrategy, InstancePartitionsType instancePartitionsType) {

Review Comment:
   Right, so kept that independent



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967314202


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigUtils.java:
##########
@@ -165,9 +166,18 @@ public static TableConfig fromZNRecord(ZNRecord znRecord)
           new TypeReference<Map<InstancePartitionsType, String>>() { });
     }
 
+    Map<InstancePartitionsType, SegmentAssignmentConfig> segmentAssignmentConfigMap = null;
+    String segmentAssignmentConfigMapString = simpleFields.get(TableConfig.SEGMENT_ASSIGNMENT_CONFIG_MAP_KEY);
+    if (segmentAssignmentConfigMapString != null) {
+      segmentAssignmentConfigMap = JsonUtils.stringToObject(segmentAssignmentConfigMapString,
+          new TypeReference<>() {

Review Comment:
   fixed



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969894035


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   updated 



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>

Review Comment:
   updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969894291


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -125,8 +133,13 @@ private List<String> assignConsumingSegment(String segmentName, InstancePartitio
       return instancesAssigned;
     } else {
       // Replica-group based assignment
-
-      checkReplication(instancePartitions);
+      // TODO: Refactor check replication this for segment assignment strategy in follow up PR

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;

Review Comment:
   updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969908627


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +67,14 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;

Review Comment:
   ohh, I missed the _assignmentStrategyMap, done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969857599


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java:
##########
@@ -388,4 +395,89 @@ public Map<String, Map<String, String>> getNonTierSegmentAssignment() {
       return _nonTierSegmentAssignment;
     }
   }
+
+  /**
+   * Returns a partition id for offline table
+   */
+  public static int getOfflineSegmentPartitionId(String segmentName, String offlineTableName,
+      HelixManager helixManager, @Nullable String partitionColumn) {
+    SegmentZKMetadata segmentZKMetadata =
+        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), offlineTableName, segmentName);
+    Preconditions.checkState(segmentZKMetadata != null,
+        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, offlineTableName);
+    return getPartitionId(segmentZKMetadata, offlineTableName, partitionColumn);
+  }
+
+  private static int getPartitionId(SegmentZKMetadata segmentZKMetadata,
+      String offlineTableName, @Nullable String partitionColumn) {

Review Comment:
   Yes, the referencing function is getOfflineSegmentPartitionId 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967315539


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignment.java:
##########
@@ -31,20 +31,17 @@
 
 /**
  * Interface for segment assignment and table rebalance.
- * <p>
- * TODO: Add SegmentAssignmentStrategy interface and support custom segment assignment strategy (e.g. cost based segment
- *       assignment). SegmentAssignmentStrategy should not be coupled with SegmentAssignment, and SegmentAssignment
- *       should be able to choose the segment assignment strategy based on the configuration.
  */
 public interface SegmentAssignment {
 
   /**
    * Initializes the segment assignment.
-   *
    * @param helixManager Helix manager
    * @param tableConfig Table config
+   * @param instancePartitionsMap Map from type (OFFLINE|CONSUMING|COMPLETED) to instance partitions
    */
-  void init(HelixManager helixManager, TableConfig tableConfig);
+  void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap);

Review Comment:
   > ```
   > 2022-09-08T23:47:33.2814091Z org.apache.pinot.controller.api.exception.TableAlreadyExistsException: Table config for testTable_OFFLINE already exists. If this is unexpected, try deleting the table to remove all metadata associated with it.
   > 2022-09-08T23:47:33.2815139Z 	at org.apache.pinot.controller.helix.core.PinotHelixResourceManager.addTable(PinotHelixResourceManager.java:1412)
   > 2022-09-08T23:47:33.2816081Z 	at org.apache.pinot.controller.helix.core.rebalance.TableRebalancerClusterStatelessTest.testRebalanceWithTiers(TableRebalancerClusterStatelessTest.java:343)
   > 2022-09-08T23:47:33.2816919Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
   > 2022-09-08T23:47:33.2817485Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
   > 2022-09-08T23:47:33.2818120Z 	at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
   > 2022-09-08T23:47:33.2818640Z 	at java.base/java.lang.reflect.Method.invoke(Method.java:566)
   > 2022-09-08T23:47:33.2819129Z 	at org.testng.internal.MethodInvocationHelper.invokeMethod(MethodInvocationHelper.java:108)
   > 2022-09-08T23:47:33.2819624Z 	at org.testng.internal.Invoker.invokeMethod(Invoker.java:661)
   > 2022-09-08T23:47:33.2820029Z 	at org.testng.internal.Invoker.invokeTestMethod(Invoker.java:869)
   > 2022-09-08T23:47:33.2820469Z 	at org.testng.internal.Invoker.invokeTestMethods(Invoker.java:1193)
   > 2022-09-08T23:47:33.2820970Z 	at org.testng.internal.TestMethodWorker.invokeTestMethods(TestMethodWorker.java:126)
   > 2022-09-08T23:47:33.2835876Z 	at org.testng.internal.TestMethodWorker.run(TestMethodWorker.java:109)
   > 2022-09-08T23:47:33.2836443Z 	at org.testng.TestRunner.privateRun(TestRunner.java:744)
   > 2022-09-08T23:47:33.2836810Z 	at org.testng.TestRunner.run(TestRunner.java:602)
   > 2022-09-08T23:47:33.2837169Z 	at org.testng.SuiteRunner.runTest(SuiteRunner.java:380)
   > 2022-09-08T23:47:33.2837552Z 	at org.testng.SuiteRunner.runSequentially(SuiteRunner.java:375)
   > 2022-09-08T23:47:33.2837963Z 	at org.testng.SuiteRunner.privateRun(SuiteRunner.java:340)
   > 2022-09-08T23:47:33.2838328Z 	at org.testng.SuiteRunner.run(SuiteRunner.java:289)
   > 2022-09-08T23:47:33.2838718Z 	at org.testng.SuiteRunnerWorker.runSuite(SuiteRunnerWorker.java:52)
   > 2022-09-08T23:47:33.2839140Z 	at org.testng.SuiteRunnerWorker.run(SuiteRunnerWorker.java:86)
   > 2022-09-08T23:47:33.2839547Z 	at org.testng.TestNG.runSuitesSequentially(TestNG.java:1301)
   > 2022-09-08T23:47:33.2839920Z 	at org.testng.TestNG.runSuitesLocally(TestNG.java:1226)
   > 2022-09-08T23:47:33.2840271Z 	at org.testng.TestNG.runSuites(TestNG.java:1144)
   > 2022-09-08T23:47:33.2840588Z 	at org.testng.TestNG.run(TestNG.java:1115)
   > 2022-09-08T23:47:33.2841002Z 	at org.apache.maven.surefire.testng.TestNGExecutor.run(TestNGExecutor.java:284)
   > 2022-09-08T23:47:33.2841558Z 	at org.apache.maven.surefire.testng.TestNGXmlTestSuite.execute(TestNGXmlTestSuite.java:75)
   > 2022-09-08T23:47:33.2842129Z 	at org.apache.maven.surefire.testng.TestNGProvider.invoke(TestNGProvider.java:119)
   > 2022-09-08T23:47:33.2842887Z 	at org.apache.maven.surefire.booter.ForkedBooter.runSuitesInProcess(ForkedBooter.java:428)
   > 2022-09-08T23:47:33.2843451Z 	at org.apache.maven.surefire.booter.ForkedBooter.execute(ForkedBooter.java:162)
   > 2022-09-08T23:47:33.2843953Z 	at org.apache.maven.surefire.booter.ForkedBooter.run(ForkedBooter.java:562)
   > 2022-09-08T23:47:33.2844449Z 	at org.apache.maven.surefire.booter.ForkedBooter.main(ForkedBooter.java:548)
   > 2022-09-08T23:47:33.2844718Z 
   > 2022-09-08T23:47:33.2845206Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testUpdateBrokerResource  Time elapsed: 0.216 s
   > 2022-09-08T23:47:33.2846308Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateDimTableTenantConfig  Time elapsed: 0.072 s
   > 2022-09-08T23:47:33.2847330Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateTenantConfig  Time elapsed: 0.052 s
   > 2022-09-08T23:47:33.2848188Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testBrokerTenant  Time elapsed: 0.953 s
   > 2022-09-08T23:47:33.2848940Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testEmptyServerTenant  Time elapsed: 0.02 s
   > 2022-09-08T23:47:33.2849714Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testServerTenant  Time elapsed: 0.388 s
   > 2022-09-08T23:47:33.2850500Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testDualModeController  Time elapsed: 9.822 s
   > 2022-09-08T23:47:33.2851344Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testHelixOnlyController  Time elapsed: 1.718 s
   > 2022-09-08T23:47:33.2852184Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testPinotOnlyController  Time elapsed: 5.248 s
   > 2022-09-08T23:47:33.2853009Z [ERROR] org.apache.pinot.controller.api.PinotBrokerRestletResourceStatelessTest.testGetBrokers  Time elapsed: 14.044 s
   > 2022-09-08T23:47:33.2853853Z [ERROR] org.apache.pinot.controller.api.PinotIngestionRestletResourceStatelessTest.testIngestEndpoint  Time elapsed: 0.492 s
   > 2022-09-08T23:47:33.2855076Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testDefaultPinotTaskManagerNoScheduler  Time elapsed: 3.905 s
   > 2022-09-08T23:47:33.2856187Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithRestart  Time elapsed: 8.857 s
   > 2022-09-08T23:47:33.2857313Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithUpdate  Time elapsed: 2.65 s
   > 2022-09-08T23:47:33.2858407Z [ERROR] org.apache.pinot.controller.validation.ValidationManagerStatelessTest.testRebuildBrokerResourceWhenBrokerAdded  Time elapsed: 0.04 s
   > 2022-09-08T23:47:33.3510290Z [INFO] 
   > 2022-09-08T23:47:33.3511030Z [INFO] Results:
   > 2022-09-08T23:47:33.3512144Z [INFO] 
   > 2022-09-08T23:47:33.3512504Z [ERROR] Failures: 
   > 2022-09-08T23:47:33.3513121Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalance:210 expected [DONE] but found [FAILED]
   > 2022-09-08T23:47:33.3534306Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalanceWithTiers:343 » TableAlreadyExists
   > ```
   > 
   > It looks that rebalancing is failing.
   
   Fixed this as well.I think this was not showing up because of other flaky test.  In the existing code, we set BalanceNumSegmentAssignmentStrategy as default and update table config replication which did not match the assignment strategy which caused this test to fail. Handled change in strategyAssignmentFactory



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967320784


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -44,6 +40,13 @@
  */
 public class OfflineSegmentAssignment extends BaseSegmentAssignment {

Review Comment:
   InstanceVariable as in? The comparison would be still same correct? 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r966612919


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -85,70 +89,41 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
       _logger.info("Initialized with replication: {} and partition column: {} for table: {}", _replication,
           _partitionColumn, _tableNameWithType);
     }
+    setSegmentAssignmentStrategyMap(instancePartitionsMap);
   }
 
   /**
-   * Returns the replication of the table.
+   * Sets Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected abstract int getReplication(TableConfig tableConfig);
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+      _assignmentStrategyMap = SegmentAssignmentStrategyFactory.getSegmentAssignmentStrategy(_helixManager,
+          _tableConfig, instancePartitionsMap);
+  }
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Returns the replication of the table.
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
+  protected abstract int getReplication(TableConfig tableConfig);
 
   /**
    * Helper method to assign instances based on the current assignment and instance partitions.
    */
   protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,

Review Comment:
   I think that we can remove this helper method and directly call the following line in `OfflineSegmentAssignment/RealtimeSegmentAssignment`
   ```
   segmentAssignmentStrategy.assignSegment(segmentName,
           currentAssignment, instancePartitions, instancePartitionsType);
   ```



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignment.java:
##########
@@ -31,20 +31,17 @@
 
 /**
  * Interface for segment assignment and table rebalance.
- * <p>
- * TODO: Add SegmentAssignmentStrategy interface and support custom segment assignment strategy (e.g. cost based segment
- *       assignment). SegmentAssignmentStrategy should not be coupled with SegmentAssignment, and SegmentAssignment
- *       should be able to choose the segment assignment strategy based on the configuration.
  */
 public interface SegmentAssignment {
 
   /**
    * Initializes the segment assignment.
-   *
    * @param helixManager Helix manager
    * @param tableConfig Table config
+   * @param instancePartitionsMap Map from type (OFFLINE|CONSUMING|COMPLETED) to instance partitions
    */
-  void init(HelixManager helixManager, TableConfig tableConfig);
+  void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap);

Review Comment:
   I don't think that we need to add `instancePartitionsMap` here if we initialize the segment assignment once in the beginning of
   
   1. assignSegment()
   
   2. rebalanceTable()
   
   and pass the object around.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   I feel that this logic becomes too complex if we also need to handle all of `OFFLINE/CONSUMING/COMPLETED`. I think that we should consider keeping this factory to be very thin (given an `AssignmentStrategy` enum, return the correct object).
   
   Instead, we should think about putting the logic for picking the correct assignment in each `OfflineSegmentAssignment / RealtimeSegmentAssignment`.
   
   This will make this logic less complex because we just need to handle `OFFLINE` from `OfflineSegmentAssignment` and `CONSUMING/COMPLETED` from `RealtimeSegmentAssignment`.
   
   
   
   



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -183,8 +176,10 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
     // Rebalance tiers first
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.COMPLETED);

Review Comment:
   Why don't we initialize the segmentAssignmentStrategy directly from `SegmentAssignmentStrategyFactory` here?
   
   In that case, we won't need to change the interface for `SegmentAssignment:init()` to take the `instancePartitionsMap`.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -194,60 +171,16 @@ protected Map<String, Map<String, String>> reassignSegments(String instanceParti
       // When bootstrap is enabled, start with an empty assignment and reassign all segments
       newAssignment = new TreeMap<>();
       for (String segment : currentAssignment.keySet()) {
-        List<String> assignedInstances = assignSegment(segment, newAssignment, instancePartitions);
+        List<String> assignedInstances = assignSegment(segment, newAssignment, instancePartitions,

Review Comment:
   This can be simply
   ```
   List<String> assignedInstances = assignmentStrategy.assignSegment(...)
   ```



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   Do we access this from anywhere? For `setSegmentAssignmentStrategyMap()`, we can push the local `tableConfig`?



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r961189561


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,67 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+
+
+/**
+ * Interface for segment assignment strategies
+ */
+public interface SegmentAssignmentStrategy {
+
+  /**
+   * Initializes the segment assignment strategy.
+   *
+   * @param helixManager Helix manager
+   * @param tableConfig Table config
+   */
+  void init(HelixManager helixManager, TableConfig tableConfig);
+
+  /**
+   * Assigns segment to instances. The assignment strategy will be configured in
+   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
+   * assignment strategy, this function will be called to assign a new segment
+   *
+   * @param segmentName Name of the segment to be assigned
+   * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
+   * @param instancePartitions Instance partitions
+   * @return List of instances to assign the segment to
+   */
+  List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType);
+
+  /**
+   * Re-assigns segment to instances. The assignment strategy will be configured in
+   * OfflineSegmentAssignment and RealtimeSegmentAssignment classes and depending on type of
+   * assignment strategy, this function will be called to re-assign a segment
+   * when the InstancePartitions has been changed.
+   *
+   * @param currentAssignment Current segment assignment of the table (map from segment name to instance state map)
+   * @param instancePartitions Instance partitions
+   * @return

Review Comment:
   `@return  newAssignment New segment assignment of the table (map from segment name to instance state map)`?



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r961189143


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,135 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
+
+  private static HelixManager _helixManager;
+  private static String _tableName;
+  private static String _partitionColumn;
+  private int _replication;
+  private TableConfig _tableConfig;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _helixManager = helixManager;
+    _tableConfig = tableConfig;
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    ReplicaGroupStrategyConfig replicaGroupStrategyConfig =
+        tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();
+    _partitionColumn = replicaGroupStrategyConfig != null ? replicaGroupStrategyConfig.getPartitionColumn() : null;
+    if (_partitionColumn == null) {
+      LOGGER.info("Initialized OfflineReplicaGroupSegmentAssignmentStrategy "

Review Comment:
   Same, do we need Offline prefix?



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r961188846


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized OfflineBalanceNumSegmentAssignmentStrategy for table: "

Review Comment:
   `OfflineBalanceNumSegmentAssignmentStrategy` or `BalanceNumSegmentAssignmentStrategy`?



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r964221907


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);

Review Comment:
   ```suggestion
       InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
       Preconditions.checkState(
           !(segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) && instancePartitions != null,
           "Failed to find OFFLINE instance partitions for table: %s", _tableNameWithType);
   ```
   One way of rewriting to not bypass the logging below and avoid calling `segmentAssignmentStrategy.reassignSegments` directly, or move `instancePartitions != null` into `validateSegmentAssignmentStrategy`.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969894943


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,110 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    String assignmentStrategy = null;
+    Map<InstancePartitionsType, SegmentAssignmentStrategy> segmentAssignmentStrategyMap = new HashMap<>();
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    InstancePartitionsType instancePartitionsType = currentTableType == TableType.OFFLINE
+        ? InstancePartitionsType.OFFLINE
+        : InstancePartitionsType.COMPLETED;
+
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      segmentAssignmentStrategyMap.put(instancePartitionsType, segmentAssignmentStrategy);

Review Comment:
   done



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969963584


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   yeah i thought this was left over line after clean up but we need to pass the table config by looking at the current code.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969963584


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   yeah i thought this was left over line after clean up but we need to pass the table config.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] klsince commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
klsince commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r1018496195


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   hi @GSharayu , I'm testing locally with RealtimeQuickStart customized a little bit, and I came into NPE issue like below. Basically, in my local test setup, the tableConfig of airlineStats_`REALTIME` has some tierConfigs; and SegmentRelocator is enabled to run periodically to move segments around based on data age as set in tierConfigs. 
   
   As suggested above, I feel it'd be safer to init segmentAssignmentStrategy with a default policy. 
   
   ```
   2022/11/09 14:43:09.474 INFO [RealtimeSegmentAssignment] [async-task-thread-1] Rebalancing tier: myTier for table: airlineStats_REALTIME with bootstrap: false, instance partitions: {"instancePartitionsName":"airlineStats
   _myTier","partitionToInstancesMap":{"0_0":["Server_192.168.0.106_7000"]}}
   2022/11/09 14:43:09.475 WARN [TableRebalancer] [async-task-thread-1] Caught exception while calculating target assignment for table: airlineStats_REALTIME, aborting the rebalance
   java.lang.NullPointerException: null
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.reassignSegments(BaseSegmentAssignment.java:158) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.BaseSegmentAssignment.rebalanceTiers(BaseSegmentAssignment.java:129) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.assignment.segment.RealtimeSegmentAssignment.rebalanceTable(RealtimeSegmentAssignment.java:201) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.rebalance.TableRebalancer.rebalance(TableRebalancer.java:219) ~[classes/:?]
           at org.apache.pinot.controller.helix.core.relocation.SegmentRelocator.lambda$processTable$0(SegmentRelocator.java:137) ~[classes/:?]
           at java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:515) [?:?]
           at java.util.concurrent.FutureTask.run(FutureTask.java:264) [?:?]
           at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128) [?:?]
           at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628) [?:?]
           at java.lang.Thread.run(Thread.java:829) [?:?]
   ```
   
   cc @npawar for awareness ^



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970324574


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -93,62 +96,22 @@ public void init(HelixManager helixManager, TableConfig tableConfig) {
   protected abstract int getReplication(TableConfig tableConfig);
 
   /**
-   * Helper method to check whether the number of replica-groups matches the table replication for replica-group based
-   * instance partitions. Log a warning if they do not match and use the one inside the instance partitions. The
-   * mismatch can happen when table is not configured correctly (table replication and numReplicaGroups does not match
-   * or replication changed without reassigning instances).
+   * Set Segment assignment strategy for different instance partitions and puts into a map of
+   * Map<InstancePartitionsType, SegmentAssignmentStrategy>
    */
-  protected void checkReplication(InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    if (numReplicaGroups != _replication) {
-      _logger.warn(
-          "Number of replica-groups in instance partitions {}: {} does not match replication in table config: {} for "
-              + "table: {}, using: {}", instancePartitions.getInstancePartitionsName(), numReplicaGroups, _replication,
-          _tableNameWithType, numReplicaGroups);
-    }
-  }
-
-  /**
-   * Helper method to assign instances based on the current assignment and instance partitions.
-   */
-  protected List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
-      InstancePartitions instancePartitions) {
-    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
-    int numPartitions = instancePartitions.getNumPartitions();
-
-    if (numReplicaGroups == 1 && numPartitions == 1) {
-      // Non-replica-group based assignment
-
-      return SegmentAssignmentUtils.assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions,
-          _replication);
-    } else {
-      // Replica-group based assignment
-
-      checkReplication(instancePartitions);
-
-      int partitionId;
-      if (_partitionColumn == null || numPartitions == 1) {
-        partitionId = 0;
-      } else {
-        // Uniformly spray the segment partitions over the instance partitions
-        partitionId = getSegmentPartitionId(segmentName) % numPartitions;
-      }
-
-      return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
-    }
+  protected void setSegmentAssignmentStrategyMap(
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    _assignmentStrategyMap = SegmentAssignmentStrategyFactory
+        .getSegmentAssignmentStrategy(_helixManager, _tableConfig, instancePartitionsMap);

Review Comment:
   Yes, Initially we were setting assignment strategy in Offline and Realtime segment assignment and getting that from factory and the naming was misaligned. We have refactored that part of code. So naming will now not be conflicting



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967316532


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   I think both of them share common code for determining replication/ partition information. This would have duplicate code for generating strategy object in both Offline and RealtimeSegmentAssignment



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r1018573095


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   I see the issue. Let me think about the fix and create a PR



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r1018579113


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java:
##########
@@ -182,9 +185,21 @@ public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, S
     boolean bootstrap =
         config.getBoolean(RebalanceConfigConstants.BOOTSTRAP, RebalanceConfigConstants.DEFAULT_BOOTSTRAP);
 
+    // TODO: remove this check after we also refactor consuming segments assignment strategy
+    // See https://github.com/apache/pinot/issues/9047
+    SegmentAssignmentStrategy segmentAssignmentStrategy = null;
+    if (completedInstancePartitions != null) {

Review Comment:
   https://github.com/apache/pinot/pull/9774
   



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee merged pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee merged PR #9309:
URL: https://github.com/apache/pinot/pull/9309


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r980773529


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,107 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {

Review Comment:
   Also, I am not expecting InstancePartitions null in the assignment factory. we call assignment strategy factory only if the instancePartitions is not null



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,87 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableNameWithType;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableNameWithType = tableConfig.getTableName();
+    SegmentsValidationAndRetentionConfig validationAndRetentionConfig = tableConfig.getValidationConfig();
+    Preconditions.checkState(validationAndRetentionConfig != null, "Validation Config is null");
+    _replication = validationAndRetentionConfig.getReplicationNumber();
+    LOGGER.info("Initialized BalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableNameWithType, _replication);

Review Comment:
   done



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,107 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy

Review Comment:
   done



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on PR #9309:
URL: https://github.com/apache/pinot/pull/9309#issuecomment-1242269212

   > ```
   > 2022-09-08T23:47:33.2814091Z org.apache.pinot.controller.api.exception.TableAlreadyExistsException: Table config for testTable_OFFLINE already exists. If this is unexpected, try deleting the table to remove all metadata associated with it.
   > 2022-09-08T23:47:33.2815139Z 	at org.apache.pinot.controller.helix.core.PinotHelixResourceManager.addTable(PinotHelixResourceManager.java:1412)
   > 2022-09-08T23:47:33.2816081Z 	at org.apache.pinot.controller.helix.core.rebalance.TableRebalancerClusterStatelessTest.testRebalanceWithTiers(TableRebalancerClusterStatelessTest.java:343)
   > 2022-09-08T23:47:33.2816919Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
   > 2022-09-08T23:47:33.2817485Z 	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
   > 2022-09-08T23:47:33.2818120Z 	at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
   > 2022-09-08T23:47:33.2818640Z 	at java.base/java.lang.reflect.Method.invoke(Method.java:566)
   > 2022-09-08T23:47:33.2819129Z 	at org.testng.internal.MethodInvocationHelper.invokeMethod(MethodInvocationHelper.java:108)
   > 2022-09-08T23:47:33.2819624Z 	at org.testng.internal.Invoker.invokeMethod(Invoker.java:661)
   > 2022-09-08T23:47:33.2820029Z 	at org.testng.internal.Invoker.invokeTestMethod(Invoker.java:869)
   > 2022-09-08T23:47:33.2820469Z 	at org.testng.internal.Invoker.invokeTestMethods(Invoker.java:1193)
   > 2022-09-08T23:47:33.2820970Z 	at org.testng.internal.TestMethodWorker.invokeTestMethods(TestMethodWorker.java:126)
   > 2022-09-08T23:47:33.2835876Z 	at org.testng.internal.TestMethodWorker.run(TestMethodWorker.java:109)
   > 2022-09-08T23:47:33.2836443Z 	at org.testng.TestRunner.privateRun(TestRunner.java:744)
   > 2022-09-08T23:47:33.2836810Z 	at org.testng.TestRunner.run(TestRunner.java:602)
   > 2022-09-08T23:47:33.2837169Z 	at org.testng.SuiteRunner.runTest(SuiteRunner.java:380)
   > 2022-09-08T23:47:33.2837552Z 	at org.testng.SuiteRunner.runSequentially(SuiteRunner.java:375)
   > 2022-09-08T23:47:33.2837963Z 	at org.testng.SuiteRunner.privateRun(SuiteRunner.java:340)
   > 2022-09-08T23:47:33.2838328Z 	at org.testng.SuiteRunner.run(SuiteRunner.java:289)
   > 2022-09-08T23:47:33.2838718Z 	at org.testng.SuiteRunnerWorker.runSuite(SuiteRunnerWorker.java:52)
   > 2022-09-08T23:47:33.2839140Z 	at org.testng.SuiteRunnerWorker.run(SuiteRunnerWorker.java:86)
   > 2022-09-08T23:47:33.2839547Z 	at org.testng.TestNG.runSuitesSequentially(TestNG.java:1301)
   > 2022-09-08T23:47:33.2839920Z 	at org.testng.TestNG.runSuitesLocally(TestNG.java:1226)
   > 2022-09-08T23:47:33.2840271Z 	at org.testng.TestNG.runSuites(TestNG.java:1144)
   > 2022-09-08T23:47:33.2840588Z 	at org.testng.TestNG.run(TestNG.java:1115)
   > 2022-09-08T23:47:33.2841002Z 	at org.apache.maven.surefire.testng.TestNGExecutor.run(TestNGExecutor.java:284)
   > 2022-09-08T23:47:33.2841558Z 	at org.apache.maven.surefire.testng.TestNGXmlTestSuite.execute(TestNGXmlTestSuite.java:75)
   > 2022-09-08T23:47:33.2842129Z 	at org.apache.maven.surefire.testng.TestNGProvider.invoke(TestNGProvider.java:119)
   > 2022-09-08T23:47:33.2842887Z 	at org.apache.maven.surefire.booter.ForkedBooter.runSuitesInProcess(ForkedBooter.java:428)
   > 2022-09-08T23:47:33.2843451Z 	at org.apache.maven.surefire.booter.ForkedBooter.execute(ForkedBooter.java:162)
   > 2022-09-08T23:47:33.2843953Z 	at org.apache.maven.surefire.booter.ForkedBooter.run(ForkedBooter.java:562)
   > 2022-09-08T23:47:33.2844449Z 	at org.apache.maven.surefire.booter.ForkedBooter.main(ForkedBooter.java:548)
   > 2022-09-08T23:47:33.2844718Z 
   > 2022-09-08T23:47:33.2845206Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testUpdateBrokerResource  Time elapsed: 0.216 s
   > 2022-09-08T23:47:33.2846308Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateDimTableTenantConfig  Time elapsed: 0.072 s
   > 2022-09-08T23:47:33.2847330Z [ERROR] org.apache.pinot.controller.helix.core.PinotHelixResourceManagerStatelessTest.testValidateTenantConfig  Time elapsed: 0.052 s
   > 2022-09-08T23:47:33.2848188Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testBrokerTenant  Time elapsed: 0.953 s
   > 2022-09-08T23:47:33.2848940Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testEmptyServerTenant  Time elapsed: 0.02 s
   > 2022-09-08T23:47:33.2849714Z [ERROR] org.apache.pinot.controller.helix.ControllerTenantStatelessTest.testServerTenant  Time elapsed: 0.388 s
   > 2022-09-08T23:47:33.2850500Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testDualModeController  Time elapsed: 9.822 s
   > 2022-09-08T23:47:33.2851344Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testHelixOnlyController  Time elapsed: 1.718 s
   > 2022-09-08T23:47:33.2852184Z [ERROR] org.apache.pinot.controller.helix.PinotControllerModeStatelessTest.testPinotOnlyController  Time elapsed: 5.248 s
   > 2022-09-08T23:47:33.2853009Z [ERROR] org.apache.pinot.controller.api.PinotBrokerRestletResourceStatelessTest.testGetBrokers  Time elapsed: 14.044 s
   > 2022-09-08T23:47:33.2853853Z [ERROR] org.apache.pinot.controller.api.PinotIngestionRestletResourceStatelessTest.testIngestEndpoint  Time elapsed: 0.492 s
   > 2022-09-08T23:47:33.2855076Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testDefaultPinotTaskManagerNoScheduler  Time elapsed: 3.905 s
   > 2022-09-08T23:47:33.2856187Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithRestart  Time elapsed: 8.857 s
   > 2022-09-08T23:47:33.2857313Z [ERROR] org.apache.pinot.controller.helix.core.minion.PinotTaskManagerStatelessTest.testPinotTaskManagerSchedulerWithUpdate  Time elapsed: 2.65 s
   > 2022-09-08T23:47:33.2858407Z [ERROR] org.apache.pinot.controller.validation.ValidationManagerStatelessTest.testRebuildBrokerResourceWhenBrokerAdded  Time elapsed: 0.04 s
   > 2022-09-08T23:47:33.3510290Z [INFO] 
   > 2022-09-08T23:47:33.3511030Z [INFO] Results:
   > 2022-09-08T23:47:33.3512144Z [INFO] 
   > 2022-09-08T23:47:33.3512504Z [ERROR] Failures: 
   > 2022-09-08T23:47:33.3513121Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalance:210 expected [DONE] but found [FAILED]
   > 2022-09-08T23:47:33.3534306Z [ERROR]   TableRebalancerClusterStatelessTest.testRebalanceWithTiers:343 » TableAlreadyExists
   > ```
   > 
   > It looks that rebalancing is failing.
   
   Fixed this as well.I think this was not showing up because of other flaky test. In the existing code, we set BalanceNumSegmentAssignmentStrategy as default and update table config replication which did not match the assignment strategy which caused this test to fail. Handled change in strategyAssignmentFactory


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967317197


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +68,15 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;
 
   @Override
-  public void init(HelixManager helixManager, TableConfig tableConfig) {
+  public void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
     _helixManager = helixManager;
     _tableNameWithType = tableConfig.getTableName();
+    _tableConfig = tableConfig;

Review Comment:
   we use table config in setSegmentAssignmentStrategyMap which is local. 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] Jackie-Jiang commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
Jackie-Jiang commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r967441628


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignment.java:
##########
@@ -31,20 +31,17 @@
 
 /**
  * Interface for segment assignment and table rebalance.
- * <p>
- * TODO: Add SegmentAssignmentStrategy interface and support custom segment assignment strategy (e.g. cost based segment
- *       assignment). SegmentAssignmentStrategy should not be coupled with SegmentAssignment, and SegmentAssignment
- *       should be able to choose the segment assignment strategy based on the configuration.
  */
 public interface SegmentAssignment {
 
   /**
    * Initializes the segment assignment.
-   *
    * @param helixManager Helix manager
    * @param tableConfig Table config
+   * @param instancePartitionsMap Map from type (OFFLINE|CONSUMING|COMPLETED) to instance partitions
    */
-  void init(HelixManager helixManager, TableConfig tableConfig);
+  void init(HelixManager helixManager, TableConfig tableConfig,
+      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap);

Review Comment:
   +1. We shouldn't pass it into the assignment



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,114 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static Map<InstancePartitionsType, SegmentAssignmentStrategy>
+  getSegmentAssignmentStrategy(HelixManager helixManager,

Review Comment:
   We don't want to pre-compute the strategy for all types. Instead, we can simplify it to only return the strategy for one type/tier: `SegmentAssignmentStrategy getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig, InstancePartitionsType instancePartitionsType, InstancePartitions instancePartitions)`



##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -89,6 +91,7 @@ public class TableConfig extends BaseJsonConfig {
   @JsonPropertyDescription(value = "Point to an existing instance partitions")
   private Map<InstancePartitionsType, String> _instancePartitionsMap;
 
+  private Map<InstancePartitionsType, SegmentAssignmentConfig> _segmentAssignmentConfigMap;

Review Comment:
   Since the key is `InstancePartitionsType`, we won't be able to configure the assignment strategy for the tiers. Consider changing the key to `String` to be able to represent tiers



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r968714839


##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -89,6 +91,7 @@ public class TableConfig extends BaseJsonConfig {
   @JsonPropertyDescription(value = "Point to an existing instance partitions")
   private Map<InstancePartitionsType, String> _instancePartitionsMap;
 
+  private Map<InstancePartitionsType, SegmentAssignmentConfig> _segmentAssignmentConfigMap;

Review Comment:
   updated 



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969962134


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,109 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, InstancePartitions instancePartitions) {

Review Comment:
   Don't we need to pass the `instanceType` as well? Even if we are not handling `CONSUMING` with the PR, let's add the parameter.



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969894590


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();

Review Comment:
   updated



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized BalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableName, _replication);
+  }
+
+  @Override
+  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    validateSegmentAssignmentStrategy(instancePartitions);
+    return SegmentAssignmentUtils
+        .assignSegmentWithoutReplicaGroup(currentAssignment, instancePartitions, _replication);
+  }
+
+  @Override
+  public Map<String, Map<String, String>> reassignSegments(Map<String, Map<String, String>> currentAssignment,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    validateSegmentAssignmentStrategy(instancePartitions);
+    Map<String, Map<String, String>> newAssignment;
+    List<String> instances =
+        SegmentAssignmentUtils.getInstancesForNonReplicaGroupBasedAssignment(instancePartitions, _replication);
+    newAssignment = SegmentAssignmentUtils
+        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, instances, _replication);
+    return newAssignment;
+  }
+
+  private void validateSegmentAssignmentStrategy(InstancePartitions instancePartitions) {
+    int numReplicaGroups = instancePartitions.getNumReplicaGroups();
+    int numPartitions = instancePartitions.getNumPartitions();
+
+    Preconditions.checkState(numReplicaGroups == 1, "Replica groups should be 1");

Review Comment:
   done



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,135 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+class ReplicaGroupSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
+
+  private static HelixManager _helixManager;
+  private static String _tableName;
+  private static String _partitionColumn;
+  private int _replication;
+  private TableConfig _tableConfig;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _helixManager = helixManager;
+    _tableConfig = tableConfig;
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    ReplicaGroupStrategyConfig replicaGroupStrategyConfig =
+        tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();

Review Comment:
   updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969908371


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +67,14 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;

Review Comment:
   we still use this from OfflineSegmentAssignment and RealtimeSegmentAssignment assign and rebalance functions



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r969908627


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -68,11 +67,14 @@ public abstract class BaseSegmentAssignment implements SegmentAssignment {
   protected String _tableNameWithType;
   protected int _replication;
   protected String _partitionColumn;
+  protected TableConfig _tableConfig;
+  protected Map<InstancePartitionsType, SegmentAssignmentStrategy> _assignmentStrategyMap;

Review Comment:
   ohh, I missed the _assignmentStrategyMap



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r970070945


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,109 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager,
+      TableConfig tableConfig, InstancePartitions instancePartitions) {

Review Comment:
   So the idea here will be to fetch the instancePartitions for given InstancePartitionType and get assignment strategy for the same, so we dont particularly need instanceType even in consuming segments



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r971260503


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/SegmentAssignmentStrategyFactory.java:
##########
@@ -0,0 +1,105 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.apache.pinot.spi.config.table.assignment.SegmentAssignmentConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Factory for SegmentAssignmentStrategy
+ */
+public class SegmentAssignmentStrategyFactory {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentAssignmentStrategyFactory.class);
+  private SegmentAssignmentStrategyFactory() {
+  }
+
+  /**
+   * Determine Segment Assignment strategy
+   */
+  public static SegmentAssignmentStrategy
+  getSegmentAssignmentStrategy(HelixManager helixManager, TableConfig tableConfig,
+      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
+    String assignmentStrategy = null;
+
+    TableType currentTableType = tableConfig.getTableType();
+    // TODO: Handle segment assignment strategy in future for CONSUMING segments in follow up PR
+    // See https://github.com/apache/pinot/issues/9047
+    // Accommodate new changes for assignment strategy
+    Map<String, SegmentAssignmentConfig>
+        segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
+
+    if (tableConfig.isDimTable()) {
+      // Segment Assignment Strategy for DIM tables
+      Preconditions.checkState(currentTableType == TableType.OFFLINE,
+          "DIM table Segment assignment Strategy is only applicable to OfflineTables");
+      SegmentAssignmentStrategy segmentAssignmentStrategy = new DimTableSegmentAssignmentStrategy();
+      segmentAssignmentStrategy.init(helixManager, tableConfig);
+      return segmentAssignmentStrategy;
+    } else {
+      // Try to determine segment assignment strategy from table config
+      if (segmentAssignmentConfigMap != null) {
+        SegmentAssignmentConfig segmentAssignmentConfig;
+        // Use the pre defined segment assignment strategy
+        segmentAssignmentConfig = segmentAssignmentConfigMap.get(instancePartitionsType.toString().toUpperCase());
+        // Segment assignment config is only applicable to offline tables and completed segments of real time tables
+        if (segmentAssignmentConfig != null) {
+          assignmentStrategy = segmentAssignmentConfig.getAssignmentStrategy();
+        }
+      }
+    }
+
+    // Use the existing information to determine segment assignment strategy
+    SegmentAssignmentStrategy segmentAssignmentStrategy;
+    if (assignmentStrategy == null) {
+      // Calculate numReplicaGroups and numPartitions to determine segment assignment strategy
+      int numReplicaGroups = instancePartitions.getNumReplicaGroups();

Review Comment:
   We check null for instance partitions both in Realtime and offline segment assignment. I also added an extra preconditions for instance partitions to be not null in assignment strategy factory



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on PR #9309:
URL: https://github.com/apache/pinot/pull/9309#issuecomment-1233409990

   @Jackie-Jiang @snleee 


-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r965283986


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -52,63 +55,55 @@ protected int getReplication(TableConfig tableConfig) {
   @Override
   public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.assignSegment(segmentName,
+          currentAssignment, null, InstancePartitionsType.OFFLINE);
+    }
     InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
     Preconditions.checkState(instancePartitions != null, "Failed to find OFFLINE instance partitions for table: %s",
         _tableNameWithType);
     _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
         _tableNameWithType);
-    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment, instancePartitions);
+    List<String> instancesAssigned = assignSegment(segmentName, currentAssignment,
+            instancePartitions, segmentAssignmentStrategy, InstancePartitionsType.OFFLINE);
     _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
         _tableNameWithType);
     return instancesAssigned;
   }
 
-  @Override
-  protected int getSegmentPartitionId(String segmentName) {
-    SegmentZKMetadata segmentZKMetadata =
-        ZKMetadataProvider.getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName);
-    Preconditions.checkState(segmentZKMetadata != null,
-        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, _tableNameWithType);
-    return getPartitionId(segmentZKMetadata);
-  }
-
-  private int getPartitionId(SegmentZKMetadata segmentZKMetadata) {
-    String segmentName = segmentZKMetadata.getSegmentName();
-    ColumnPartitionMetadata partitionMetadata =
-        segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_partitionColumn);
-    Preconditions.checkState(partitionMetadata != null,
-        "Segment ZK metadata for segment: %s of table: %s does not contain partition metadata for column: %s",
-        segmentName, _tableNameWithType, _partitionColumn);
-    Set<Integer> partitions = partitionMetadata.getPartitions();
-    Preconditions.checkState(partitions.size() == 1,
-        "Segment ZK metadata for segment: %s of table: %s contains multiple partitions for column: %s", segmentName,
-        _tableNameWithType, _partitionColumn);
-    return partitions.iterator().next();
-  }
-
   @Override
   public Map<String, Map<String, String>> rebalanceTable(Map<String, Map<String, String>> currentAssignment,
       Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap, @Nullable List<Tier> sortedTiers,
       @Nullable Map<String, InstancePartitions> tierInstancePartitionsMap, Configuration config) {
     InstancePartitions offlineInstancePartitions = instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
+    SegmentAssignmentStrategy segmentAssignmentStrategy = _assignmentStrategyMap.get(InstancePartitionsType.OFFLINE);
+    // Need to check assignmentStrategy for Dim Tables as following pre conditions state would return false
+    if (segmentAssignmentStrategy instanceof DimTableSegmentAssignmentStrategy) {
+      return segmentAssignmentStrategy.reassignSegments(currentAssignment,
+          offlineInstancePartitions, InstancePartitionsType.OFFLINE);

Review Comment:
   For Dim tables, we would have reassignSegments for all servers, so bootstrapping and this is the same



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] GSharayu commented on a diff in pull request #9309: Introduce Segment AssignmentStrategy Interface

Posted by GitBox <gi...@apache.org>.
GSharayu commented on code in PR #9309:
URL: https://github.com/apache/pinot/pull/9309#discussion_r965280518


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized OfflineBalanceNumSegmentAssignmentStrategy for table: "
+            + "{} with replication: {}", _tableName, _tableName, _replication);

Review Comment:
   done



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/BalanceNumSegmentAssignmentStrategy.java:
##########
@@ -0,0 +1,82 @@
+/**
+ * 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.controller.helix.core.assignment.segment.strategy;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Balance num Segment assignment strategy class for offline segment assignment
+ * <ul>
+ *   <li>
+ *     <p>This segment assignment strategy is used when table replication/ num_replica_groups = 1.</p>
+ *   </li>
+ * </ul>
+ */
+public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
+  private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
+
+  private String _tableName;
+  private int _replication;
+
+  @Override
+  public void init(HelixManager helixManager, TableConfig tableConfig) {
+    _tableName = tableConfig.getTableName();
+    _replication = tableConfig.getValidationConfig().getReplicationNumber();
+    LOGGER.info("Initialized OfflineBalanceNumSegmentAssignmentStrategy for table: "

Review Comment:
   updated



-- 
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-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org