@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 testBootstrapTableWithPartition() {
    Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
    for (String segmentName : SEGMENTS) {
      List<String> instancesAssigned = _segmentAssignmentWithPartition
          .assignSegment(segmentName, currentAssignment, _instancePartitionsMapWithPartition);
      currentAssignment
          .put(segmentName, SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned, SegmentStateModel.ONLINE));
    }

    // Bootstrap table should reassign all segments based on their alphabetical order within the partition
    RebalanceConfig rebalanceConfig = new RebalanceConfig();
    rebalanceConfig.setBootstrap(true);
    Map<String, Map<String, String>> newAssignment =
        _segmentAssignmentWithPartition.rebalanceTable(currentAssignment, _instancePartitionsMapWithPartition, null,
            null, rebalanceConfig);
    assertEquals(newAssignment.size(), NUM_SEGMENTS);
    int numSegmentsPerPartition = NUM_SEGMENTS / NUM_PARTITIONS;
    String[][] partitionIdToSegmentsMap = new String[NUM_PARTITIONS][numSegmentsPerPartition];
    for (int i = 0; i < NUM_SEGMENTS; i++) {
      partitionIdToSegmentsMap[i % NUM_PARTITIONS][i / NUM_PARTITIONS] = SEGMENTS.get(i);
    }
    String[][] partitionIdToSortedSegmentsMap = new String[NUM_PARTITIONS][numSegmentsPerPartition];
    for (int i = 0; i < NUM_PARTITIONS; i++) {
      String[] sortedSegments = new String[numSegmentsPerPartition];
      System.arraycopy(partitionIdToSegmentsMap[i], 0, sortedSegments, 0, numSegmentsPerPartition);
      Arrays.sort(sortedSegments);
      partitionIdToSortedSegmentsMap[i] = sortedSegments;
    }
    for (int i = 0; i < NUM_PARTITIONS; i++) {
      for (int j = 0; j < numSegmentsPerPartition; j++) {
        assertEquals(newAssignment.get(partitionIdToSortedSegmentsMap[i][j]),
            currentAssignment.get(partitionIdToSegmentsMap[i][j]));
      }
    }
  }