[GitHub] spark pull request #20146: [SPARK-11215][ML] Add multiple columns support to...

[holdenk <gi...@git.apache.org>]

Github user holdenk commented on a diff in the pull request:

    https://github.com/apache/spark/pull/20146#discussion_r239887472
  
    --- Diff: mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala ---
    @@ -130,21 +159,60 @@ class StringIndexer @Since("1.4.0") (
       @Since("1.4.0")
       def setOutputCol(value: String): this.type = set(outputCol, value)
     
    +  /** @group setParam */
    +  @Since("2.4.0")
    +  def setInputCols(value: Array[String]): this.type = set(inputCols, value)
    +
    +  /** @group setParam */
    +  @Since("2.4.0")
    +  def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
    +
    +  private def countByValue(
    +      dataset: Dataset[_],
    +      inputCols: Array[String]): Array[OpenHashMap[String, Long]] = {
    +
    +    val aggregator = new StringIndexerAggregator(inputCols.length)
    +    implicit val encoder = Encoders.kryo[Array[OpenHashMap[String, Long]]]
    +
    +    dataset.select(inputCols.map(col(_).cast(StringType)): _*)
    +      .toDF
    +      .groupBy().agg(aggregator.toColumn)
    +      .as[Array[OpenHashMap[String, Long]]]
    +      .collect()(0)
    +  }
    +
       @Since("2.0.0")
       override def fit(dataset: Dataset[_]): StringIndexerModel = {
         transformSchema(dataset.schema, logging = true)
    -    val values = dataset.na.drop(Array($(inputCol)))
    -      .select(col($(inputCol)).cast(StringType))
    -      .rdd.map(_.getString(0))
    -    val labels = $(stringOrderType) match {
    -      case StringIndexer.frequencyDesc => values.countByValue().toSeq.sortBy(-_._2)
    -        .map(_._1).toArray
    -      case StringIndexer.frequencyAsc => values.countByValue().toSeq.sortBy(_._2)
    -        .map(_._1).toArray
    -      case StringIndexer.alphabetDesc => values.distinct.collect.sortWith(_ > _)
    -      case StringIndexer.alphabetAsc => values.distinct.collect.sortWith(_ < _)
    -    }
    -    copyValues(new StringIndexerModel(uid, labels).setParent(this))
    +
    +    val (inputCols, _) = getInOutCols()
    +
    +    val filteredDF = dataset.na.drop(inputCols)
    +
    +    // In case of equal frequency when frequencyDesc/Asc, we further sort the strings by alphabet.
    --- End diff --
    
    Is this a change of behaviour? If so we should document it in a visible place.


---

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