fix

Signed-off-by: Weichen Xu <[email protected]>
WeichenXu123 · Oct 10, 2024 · b16d64a · b16d64a
1 parent 8a03a3b
commit b16d64a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala b/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.ml.util
 
+import org.apache.spark.SparkIllegalArgumentException
 import org.apache.spark.ml.attribute._
 import org.apache.spark.ml.linalg.VectorUDT
 import org.apache.spark.sql.catalyst.util.AttributeNameParser
@@ -214,7 +215,11 @@ private[spark] object SchemaUtils {
    */
   def getSchemaField(schema: StructType, colName: String): StructField = {
     val colSplits = AttributeNameParser.parseAttributeName(colName)
-    schema.findNestedField(colSplits, resolver = SQLConf.get.resolver).get._2
+    val fieldOpt = schema.findNestedField(colSplits, resolver = SQLConf.get.resolver)
+    if (fieldOpt.isEmpty) {
+      throw new SparkIllegalArgumentException("FIELD_NOT_FOUND")
+    }
+    fieldOpt.get._2
   }
 
   /**