viirya commented on a change in pull request #20146: [SPARK-11215][ML] Add
multiple columns support to StringIndexer
URL: https://github.com/apache/spark/pull/20146#discussion_r243732475
##########
File path: mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
##########
@@ -257,29 +372,74 @@ class StringIndexerModel (
} else if (keepInvalid) {
labels.length
} else {
- throw new SparkException(s"Unseen label: $label. To handle unseen
labels, " +
+ throw new SparkException(s"Unseen label: $label. To handle unseen
labels, " +
s"set Param handleInvalid to ${StringIndexer.KEEP_INVALID}.")
}
}
}.asNondeterministic()
+ }
+
+ @Since("2.0.0")
+ override def transform(dataset: Dataset[_]): DataFrame = {
+ transformSchema(dataset.schema, logging = true)
+
+ var (inputColNames, outputColNames) = getInOutCols()
+ val outputColumns = new Array[Column](outputColNames.length)
+
+ // Skips invalid rows if `handleInvalid` is set to
`StringIndexer.SKIP_INVALID`.
+ val filteredDataset = if (getHandleInvalid == StringIndexer.SKIP_INVALID) {
+ filterInvalidData(dataset, inputColNames)
+ } else {
+ dataset
+ }
+
+ for (i <- 0 until outputColNames.length) {
+ val inputColName = inputColNames(i)
+ val outputColName = outputColNames(i)
+ val labelToIndex = labelsToIndexArray(i)
+ val labels = labelsArray(i)
+
+ if (!dataset.schema.fieldNames.contains(inputColName)) {
+ logWarning(s"Input column ${inputColName} does not exist during
transformation. " +
+ "Skip StringIndexerModel for this column.")
+ outputColNames(i) = null
+ } else {
+ val filteredLabels = getHandleInvalid match {
+ case StringIndexer.KEEP_INVALID => labels :+ "__unknown"
+ case _ => labels
+ }
+ val metadata = NominalAttribute.defaultAttr
+ .withName(outputColName)
+ .withValues(filteredLabels)
+ .toMetadata()
+ val keepInvalid = (getHandleInvalid == StringIndexer.KEEP_INVALID)
Review comment:
Removed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
