make sure the code works for Float type and add the unit test

Author: lu-wang-dl

mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala

@@ -129,14 +129,21 @@ class KMeansModel private[ml] (
   override def transform(dataset: Dataset[_]): DataFrame = {
     transformSchema(dataset.schema, logging = true)
     // val predictUDF = udf((vector: Vector) => predict(vector))
-    if (dataset.schema($(featuresCol)).dataType.equals(new VectorUDT)) {
-      val predictUDF = udf((vector: Vector) => predict(vector))
-      dataset.withColumn($(predictionCol), predictUDF(col($(featuresCol))))
-    } else {
-      val predictUDF = udf((vector: Seq[_]) =>
-        predict(Vectors.dense(vector.asInstanceOf[Seq[Double]].toArray)))
-      dataset.withColumn($(predictionCol), predictUDF(col($(featuresCol))))
+    val predictUDF = if (dataset.schema($(featuresCol)).dataType.equals(new VectorUDT)) {
+      udf((vector: Vector) => predict(vector))
     }
+    else {
+      udf((vector: Seq[_]) => {
+        val featureArray = Array.fill[Double](vector.size)(0.0)
+        for (idx <- 0 until vector.size) {
+          featureArray(idx) = vector(idx).toString().toDouble
+        }
+        OldVectors.fromML(Vectors.dense(featureArray))
+        predict(Vectors.dense(featureArray))
+      })
+    }
+
+    dataset.withColumn($(predictionCol), predictUDF(col($(featuresCol))))
   }
 
   @Since("1.5.0")
@@ -164,6 +171,12 @@ class KMeansModel private[ml] (
     SchemaUtils.checkColumnTypes(dataset.schema, $(featuresCol), typeCandidates)
     val data: RDD[OldVector] = dataset.select(col($(featuresCol))).rdd.map {
       case Row(point: Vector) => OldVectors.fromML(point)
+      case Row(point: Seq[_]) =>
+        val featureArray = Array.fill[Double](point.size)(0.0)
+        for (idx <- 0 until point.size) {
+          featureArray(idx) = point(idx).toString().toDouble
+        }
+        OldVectors.fromML(Vectors.dense(featureArray))
     }
     parentModel.computeCost(data)
   }
@@ -330,8 +343,12 @@ class KMeans @Since("1.5.0") (
     val instances: RDD[OldVector] = dataset.select(col($(featuresCol))).rdd.map {
       case Row(point: Vector) => OldVectors.fromML(point)
       case Row(point: Seq[_]) =>
-        OldVectors.fromML(Vectors.dense(point.asInstanceOf[Seq[Double]].toArray))
-    }
+        val featureArray = Array.fill[Double](point.size)(0.0)
+        for (idx <- 0 until point.size) {
+          featureArray(idx) = point(idx).toString().toDouble
+        }
+        OldVectors.fromML(Vectors.dense(featureArray))
+      }
 
     if (handlePersistence) {
       instances.persist(StorageLevel.MEMORY_AND_DISK)

mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala

@@ -27,6 +27,7 @@ import org.apache.spark.mllib.clustering.{DistanceMeasure, KMeans => MLlibKMeans
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
 import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types.{ArrayType, DoubleType, FloatType, IntegerType, StructType}
 
 private[clustering] case class TestRow(features: Vector)
 
@@ -196,30 +197,43 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultR
   }
 
   test("KMean with Array input") {
-    val featuresColName = "array_model_features"
+    val featuresColNameD = "array_double_features"
+    val featuresColNameF = "array_float_features"
 
-    val arrayUDF = udf { (features: Vector) =>
-      features.toArray
+    val doubleUDF = udf { (features: Vector) =>
+      val featureArray = Array.fill[Double](features.size)(0.0)
+      features.foreachActive((idx, value) => featureArray(idx) = value.toFloat)
+      featureArray
+    }
+    val floatUDF = udf { (features: Vector) =>
+      val featureArray = Array.fill[Float](features.size)(0.0f)
+      features.foreachActive((idx, value) => featureArray(idx) = value.toFloat)
+      featureArray
     }
-    val newdataset = dataset.withColumn(featuresColName, arrayUDF(col("features")) )
 
-    val kmeans = new KMeans()
-      .setFeaturesCol(featuresColName)
+    val newdatasetD = dataset.withColumn(featuresColNameD, doubleUDF(col("features")))
+      .drop("features")
+    val newdatasetF = dataset.withColumn(featuresColNameF, floatUDF(col("features")))
+      .drop("features")
 
-    assert(kmeans.getK === 2)
-    assert(kmeans.getFeaturesCol === featuresColName)
-    assert(kmeans.getPredictionCol === "prediction")
-    assert(kmeans.getMaxIter === 20)
-    assert(kmeans.getInitMode === MLlibKMeans.K_MEANS_PARALLEL)
-    assert(kmeans.getInitSteps === 2)
-    assert(kmeans.getTol === 1e-4)
-    assert(kmeans.getDistanceMeasure === DistanceMeasure.EUCLIDEAN)
-    val model = kmeans.setMaxIter(1).fit(newdataset)
+    assert(newdatasetD.schema(featuresColNameD).dataType.equals(new ArrayType(DoubleType, false)))
+    assert(newdatasetF.schema(featuresColNameF).dataType.equals(new ArrayType(FloatType, false)))
+
+    val kmeansD = new KMeans().setK(k).setFeaturesCol(featuresColNameD).setSeed(1)
+    val kmeansF = new KMeans().setK(k).setFeaturesCol(featuresColNameF).setSeed(1)
+    val modelD = kmeansD.fit(newdatasetD)
+    val modelF = kmeansF.fit(newdatasetF)
+
+    val transformedD = modelD.transform(newdatasetD)
+    val transformedF = modelF.transform(newdatasetF)
+
+    val predictDifference = transformedD.select("prediction")
+      .except(transformedF.select("prediction"))
+
+    assert(predictDifference.count() == 0)
+
+    assert(modelD.computeCost(newdatasetD) == modelF.computeCost(newdatasetF) )
 
-    MLTestingUtils.checkCopyAndUids(kmeans, model)
-    assert(model.hasSummary)
-    val copiedModel = model.copy(ParamMap.empty)
-    assert(copiedModel.hasSummary)
   }