@Override
  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
    int numPartitions = instancePartitions.getNumPartitions();
    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;
      } else {
        partitionId = SegmentAssignmentUtils
            .getRealtimeSegmentPartitionId(segmentName, _tableName, _helixManager, _partitionColumn) % numPartitions;
      }
    }
    return SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment, instancePartitions, partitionId);
  }
@Test
  public void testAssignSegmentWithPartition() {
    int numInstancesPerReplicaGroup = NUM_INSTANCES / NUM_REPLICAS;
    Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
    int numInstancesPerPartition = numInstancesPerReplicaGroup / NUM_PARTITIONS;
    for (int segmentId = 0; segmentId < NUM_SEGMENTS; segmentId++) {
      String segmentName = SEGMENTS.get(segmentId);
      List<String> instancesAssigned = _segmentAssignmentWithPartition
          .assignSegment(segmentName, currentAssignment, _instancePartitionsMapWithPartition);
      assertEquals(instancesAssigned.size(), NUM_REPLICAS);

      // Segment 0 (partition 0) should be assigned to instance 0, 6, 12
      // Segment 1 (partition 1) should be assigned to instance 2, 8, 14
      // Segment 2 (partition 2) should be assigned to instance 4, 10, 16
      // Segment 3 (partition 0) should be assigned to instance 1, 7, 13
      // Segment 4 (partition 1) should be assigned to instance 3, 9, 15
      // Segment 5 (partition 2) should be assigned to instance 5, 11, 17
      // Segment 6 (partition 0) should be assigned to instance 0, 6, 12
      // Segment 7 (partition 1) should be assigned to instance 2, 8, 14
      // ...
      int partitionId = segmentId % NUM_PARTITIONS;
      for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) {
        int expectedAssignedInstanceId =
            (segmentId % numInstancesPerReplicaGroup) / NUM_PARTITIONS + partitionId * numInstancesPerPartition
                + replicaGroupId * numInstancesPerReplicaGroup;
        assertEquals(instancesAssigned.get(replicaGroupId), INSTANCES.get(expectedAssignedInstanceId));
      }

      currentAssignment
          .put(segmentName, SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned, SegmentStateModel.ONLINE));
    }
  }