[GitHub] [spark] srowen commented on a change in pull request #20146: [SPARK-11215][ML] Add multiple columns support to StringIndexer

[GitBox <gi...@apache.org>]

srowen commented on a change in pull request #20146:
URL: https://github.com/apache/spark/pull/20146#discussion_r534479668



##########
File path: mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
##########
@@ -172,37 +265,73 @@ object StringIndexer extends DefaultParamsReadable[StringIndexer] {
 
   @Since("1.6.0")
   override def load(path: String): StringIndexer = super.load(path)
+
+  // Returns a function used to sort strings by frequency (ascending or descending).
+  // In case of equal frequency, it sorts strings by alphabet (ascending).
+  private[feature] def getSortFunc(
+      ascending: Boolean): ((String, Long), (String, Long)) => Boolean = {
+    if (ascending) {
+     (a: (String, Long), b: (String, Long)) => {
+       if (a._2 == b._2) {
+         a._1 < b._1
+       } else {
+         a._2 < b._2
+       }
+     }
+    } else {
+      (a: (String, Long), b: (String, Long)) => {
+        if (a._2 == b._2) {
+          a._1 < b._1
+        } else {
+          a._2 > b._2
+        }
+      }
+    }
+  }
 }
 
 /**
  * Model fitted by [[StringIndexer]].
  *
- * @param labels  Ordered list of labels, corresponding to indices to be assigned.
+ * @param labelsArray Array of ordered list of labels, corresponding to indices to be assigned
+ *                    for each input column.
  *
- * @note During transformation, if the input column does not exist,
- * `StringIndexerModel.transform` would return the input dataset unmodified.
+ * @note During transformation, if any input column does not exist,
+ * `StringIndexerModel.transform` would skip the input column.
+ * If all input columns do not exist, it returns the input dataset unmodified.
  * This is a temporary fix for the case when target labels do not exist during prediction.
  */
 @Since("1.4.0")
 class StringIndexerModel (
     @Since("1.4.0") override val uid: String,
-    @Since("1.5.0") val labels: Array[String])
+    @Since("3.0.0") val labelsArray: Array[Array[String]])
   extends Model[StringIndexerModel] with StringIndexerBase with MLWritable {
 
   import StringIndexerModel._
 
   @Since("1.5.0")
-  def this(labels: Array[String]) = this(Identifiable.randomUID("strIdx"), labels)
-
-  private val labelToIndex: OpenHashMap[String, Double] = {
-    val n = labels.length
-    val map = new OpenHashMap[String, Double](n)
-    var i = 0
-    while (i < n) {
-      map.update(labels(i), i)
-      i += 1
+  def this(labels: Array[String]) = this(Identifiable.randomUID("strIdx"), Array(labels))
+
+  @Since("3.0.0")
+  def this(labelsArray: Array[Array[String]]) = this(Identifiable.randomUID("strIdx"), labelsArray)
+
+  @Since("1.5.0")
+  def labels: Array[String] = {

Review comment:
       Oh @viirya I just realized that we didn't add a labelsArray method to Pyspark - looks like that is necessary right? otherwise can't access this when multiple cols are fit




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org