@Test
public void getKmeansPredictColIdxsTest() {

    KMeansTrainModelData.ParamSummary params = new KMeansTrainModelData.ParamSummary();
    params.vectorColName = "vec";
    params.latitudeColName = null;
    params.longtitudeColName = null;
    params.k = 2;
    params.vectorSize = vectorSize;
    params.distanceType = HasKMeansWithHaversineDistanceType.DistanceType.EUCLIDEAN;
    String[] colNames = new String[] {"id", "vec", "lat", "lon"};
    Assert.assertArrayEquals(KMeansUtil.getKmeansPredictColIdxs(params, colNames), new int[] {1});
    params.vectorColName = null;
    params.latitudeColName = "lat";
    params.longtitudeColName = "lon";
    params.distanceType = HasKMeansWithHaversineDistanceType.DistanceType.HAVERSINE;
    Assert.assertArrayEquals(KMeansUtil.getKmeansPredictColIdxs(params, colNames), new int[] {2, 3});
}