update

Signed-off-by: Weichen Xu <[email protected]>
WeichenXu123 · Jul 12, 2024 · f91fedf · f91fedf
1 parent 835bb3e
commit f91fedf
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -28,6 +28,7 @@ import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.{Column, DataFrame, Dataset, Encoder, Encoders, Row}
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.expressions.{If, Literal}
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.functions._
@@ -136,7 +137,7 @@ private[feature] trait StringIndexerBase extends Params with HasHandleInvalid wi
   }
 
   protected def extractInputDataType(schema: StructType, inputColName: String): Option[DataType] = {
-    val inputSplits = inputColName.split("\\.")
+    val inputSplits = UnresolvedAttribute.parseAttributeName(inputColName)
     var dtype: Option[DataType] = Some(schema)
     var i = 0
     while (i < inputSplits.length && dtype.isDefined) {

diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
@@ -150,12 +150,12 @@ class StringIndexerSuite extends MLTest with DefaultReadWriteTest {
       (5, 1.0)
     ).toDF("id", "labelIndex")
 
-    testTransformerByGlobalCheckFunc[(Int, String)](df, indexerModel, "id", "labelIndex") { rows =>
-      val attr = Attribute.fromStructField(rows.head.schema("labelIndex"))
-        .asInstanceOf[NominalAttribute]
-      assert(attr.values.get === Array("a", "c", "b"))
-      assert(rows === expected.collect().toSeq)
-    }
+    val dfOutput = indexerModel.transform(df)
+    val outputs = dfOutput.select("id", "labelIndex").collect().toSeq
+    val attr = Attribute.fromStructField(outputs.head.schema("labelIndex"))
+      .asInstanceOf[NominalAttribute]
+    assert(attr.values.get === Array("a", "c", "b"))
+    assert(outputs === expected.collect().toSeq)
   }
 
   test("StringIndexerUnseen") {