diff --git a/project/Common.scala b/project/Common.scala
index b9228fb..e0cea79 100644
--- a/project/Common.scala
+++ b/project/Common.scala
@@ -16,7 +16,7 @@ object Common {
     git.baseVersion := "0.0.1",
     parallelExecution in test := false,
     updateOptions := updateOptions.value.withCachedResolution(true),
-    sparkVersion := "3.0.1",
+    sparkVersion := "3.1.1",
     sparkComponents += "mllib",
     spIgnoreProvided := true
   )
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
index e5d9109..7819fcd 100644
--- a/project/Dependencies.scala
+++ b/project/Dependencies.scala
@@ -12,7 +12,7 @@ object Dependencies {
 
     object Test {
       val scalatest = "org.scalatest" %% "scalatest" % "3.1.0" % "test"
-      val sparktest = "org.apache.spark" %% "spark-core" % "3.0.1"  % "test" classifier "tests"
+      val sparktest = "org.apache.spark" %% "spark-core" % "3.1.1"  % "test" classifier "tests"
     }
   }
 
diff --git a/spark-knn-core/src/main/scala/org/apache/spark/ml/classification/KNNClassifier.scala b/spark-knn-core/src/main/scala/org/apache/spark/ml/classification/KNNClassifier.scala
index aedb43f..100a196 100644
--- a/spark-knn-core/src/main/scala/org/apache/spark/ml/classification/KNNClassifier.scala
+++ b/spark-knn-core/src/main/scala/org/apache/spark/ml/classification/KNNClassifier.scala
@@ -7,6 +7,7 @@ import org.apache.spark.ml.param.shared.HasWeightCol
 import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
 import org.apache.spark.ml.linalg._
 import org.apache.spark.ml.feature.LabeledPoint
+import org.apache.spark.ml.stat.MultiClassSummarizer
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.{DoubleType, StructType}
 import org.apache.spark.sql.{DataFrame, Dataset, Row}
diff --git a/spark-knn-examples/src/main/scala/org/apache/spark/ml/classification/NaiveKNN.scala b/spark-knn-examples/src/main/scala/org/apache/spark/ml/classification/NaiveKNN.scala
index e967eeb..e6a97b5 100644
--- a/spark-knn-examples/src/main/scala/org/apache/spark/ml/classification/NaiveKNN.scala
+++ b/spark-knn-examples/src/main/scala/org/apache/spark/ml/classification/NaiveKNN.scala
@@ -8,6 +8,7 @@ import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
 import org.apache.spark.ml.{Model, Predictor}
 import org.apache.spark.ml.linalg._
 import org.apache.spark.ml.feature.LabeledPoint
+import org.apache.spark.ml.stat.MultiClassSummarizer
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.{ArrayType, DoubleType, StructType}
 import org.apache.spark.sql.{DataFrame, Dataset, Row}