Jackie-Jiang commented on a change in pull request #4533: [Instance Assignment]
De-couple assignment strategy from SegmentAssignment
URL: https://github.com/apache/incubator-pinot/pull/4533#discussion_r315414688
##
File path:
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java
##
@@ -0,0 +1,260 @@
+/**
+ * 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;
+
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.TreeMap;
+import org.apache.commons.configuration.Configuration;
+import org.apache.helix.HelixManager;
+import org.apache.pinot.common.config.ReplicaGroupStrategyConfig;
+import org.apache.pinot.common.config.TableConfig;
+import org.apache.pinot.common.metadata.ZKMetadataProvider;
+import org.apache.pinot.common.metadata.segment.ColumnPartitionMetadata;
+import org.apache.pinot.common.metadata.segment.OfflineSegmentZKMetadata;
+import org.apache.pinot.common.utils.InstancePartitionsType;
+import org.apache.pinot.common.utils.Pairs;
+import org.apache.pinot.controller.helix.core.assignment.InstancePartitions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Segment assignment for offline table.
+ *
+ *
+ * Non-replica-group based assignment (only 1 replica in instance
partitions):
+ * Assign the segment to the instance with the least number of
segments. In case of a tie, assign the segment to
+ * the instance with the smallest index in the list. Use Helix
AutoRebalanceStrategy to rebalance the table.
+ *
+ *
+ * Replica-group based assignment (more than 1 replicas in instance
partitions):
+ * Among replicas, always mirror the assignment (pick the same index of
the instance).
+ * Within each partition, assign the segment to the servers with the
least segments already assigned. In case of
+ * a tie, assign to the server with the smallest index in the list. Do
this for one replica and mirror the
+ * assignment to other replicas.
+ * To rebalance a table, within each partition, first calculate the
number of segments on each server, loop over
+ * all the segments and keep the assignment if number of segments for the
server has not been reached and track the
+ * not assigned segments, then assign the left-over segments to the
servers with the least segments, or the smallest
+ * index if there is a tie. Repeat the process for all the partitions in
one replica, and mirror the assignment to
+ * other replicas. With this greedy algorithm, the result is deterministic
and with minimum segment moves.
+ *
+ *
+ */
+public class OfflineSegmentAssignment implements SegmentAssignment {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(OfflineSegmentAssignment.class);
+
+ private HelixManager _helixManager;
+ private String _offlineTableName;
+ private int _replication;
+ private String _partitionColumn;
+
+ @Override
+ public void init(HelixManager helixManager, TableConfig tableConfig) {
+_helixManager = helixManager;
+_offlineTableName = tableConfig.getTableName();
+_replication = tableConfig.getValidationConfig().getReplicationNumber();
+ReplicaGroupStrategyConfig strategyConfig =
tableConfig.getValidationConfig().getReplicaGroupStrategyConfig();
+_partitionColumn = strategyConfig != null ?
strategyConfig.getPartitionColumn() : null;
+
+if (_partitionColumn == null) {
+ LOGGER.info("Initialized OfflineSegmentAssignment with replication: {}
without partition column for table: {} ",
+ _replication, _offlineTableName);
+} else {
+ LOGGER.info("Initialized OfflineSegmentAssignment with replication: {}
and partition column: {} for table: {}",
+ _replication, _partitionColumn, _offlineTableName);
+}
+ }
+
+ @Override
+ public List assignSegment(String segmentName, Map> currentAssignment,
+ Map instancePartitionsMap) {
+I
