KMeansUtil implements Serializable { public static int[] getKmeansPredictColIdxs(KMeansTrainModelData.ParamSummary params, String[] dataCols) { Preconditions.checkArgument((null == params.longtitudeColName) == (null == params.latitudeColName), "Model Format error!"); Preconditions.checkArgument(params.distanceType.equals(HasKMeansWithHaversineDistanceType.DistanceType.HAVERSINE) == (null == params.vectorColName && null != params.longtitudeColName), "Model Format error!"); int[] colIdxs; if (null != params.vectorColName) { colIdxs = new int[1]; colIdxs[0] = TableUtil.findColIndexWithAssert(dataCols, params.vectorColName); } else { colIdxs = new int[2]; colIdxs[0] = TableUtil.findColIndexWithAssert(dataCols, params.latitudeColName); colIdxs[1] = TableUtil.findColIndexWithAssert(dataCols, params.longtitudeColName); } return colIdxs; }  static FastDistanceMatrixData buildCentroidsMatrix(List<FastDistanceVectorData> vectors,
                                                              FastDistance distance,
                                                              int vectorSize); static int updateSumMatrix(FastDistanceVectorData sample,
                                      long sampleWeight,
                                      FastDistanceMatrixData centroids,
                                      int vectorSize,
                                      double[] sumMatrix,
                                      int k,
                                      FastDistance fastDistance,
                                      DenseMatrix distanceMatrix); static Tuple2<Integer, Double> getClosestClusterIndex(FastDistanceVectorData sample,
                                                                 FastDistanceMatrixData centroids,
                                                                 int k,
                                                                 FastDistance distance,
                                                                 DenseMatrix distanceMatrix); static double[] getClusterDistances(FastDistanceVectorData sample,
                                               FastDistanceMatrixData centroids,
                                               FastDistance distance,
                                               DenseMatrix distanceMatrix); static Tuple2<Integer, Double> getClosestClusterIndex(KMeansTrainModelData trainModelData,
                                                                 Vector sample,
                                                                 ContinuousDistance distance); static double[] getClusterDistances(KMeansTrainModelData trainModelData,
                                               Vector sample,
                                               ContinuousDistance distance); static int getMinPointIndex(double[] data, int endIndex); static int[] getKmeansPredictColIdxs(KMeansTrainModelData.ParamSummary params, String[] dataCols); static Vector getKMeansPredictVector(int[] colIdxs, Row row); static KMeansTrainModelData transformPredictDataToTrainData(KMeansPredictModelData predictModelData); static KMeansPredictModelData transformTrainDataToPredictData(KMeansTrainModelData trainModelData); static double[] getProbArrayFromDistanceArray(double[] distances); static KMeansTrainModelData loadModelForTrain(Params params, Iterable<String> data);  }
@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}); }