Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/7574#discussion_r35279311
--- Diff: mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
---
@@ -130,9 +173,52 @@ class RFormula(override val uid: String)
"Label column already exists and is not of type DoubleType.")
}
- private def hasLabelCol(schema: StructType): Boolean = {
- schema.map(_.name).contains($(labelCol))
+ private def featureTransformer(schema: StructType): Transformer = {
+ // StringType terms and terms representing interactions need to be
encoded before assembly.
+ // TODO(ekl) add support for feature interactions
+ var encoderStages = Seq[Transformer]()
+ var tempColumns = Seq[String]()
+ val encodedTerms = parsedFormula.terms.map { term =>
+ schema(term) match {
+ case column if column.dataType == StringType =>
+ val encodedTerm = term + "_onehot_" + uid
+ val indexer = factorLevels(term)
+ val indexCol = indexer.getOrDefault(indexer.outputCol)
+ encoderStages :+= indexer
+ encoderStages :+= new OneHotEncoder()
+ .setInputCol(indexCol)
+ .setOutputCol(encodedTerm)
+ tempColumns :+= encodedTerm
+ tempColumns :+= indexCol
+ encodedTerm
+ case _ =>
+ term
+ }
+ }
+ encoderStages :+= new VectorAssembler(uid)
+ .setInputCols(encodedTerms.toArray)
+ .setOutputCol($(featuresCol))
+ encoderStages :+= new ColumnPruner(tempColumns.toSet)
+ new PipelineModel(uid, encoderStages.toArray)
+ }
+}
+
+/**
+ * Utility transformer for removing temporary columns from a DataFrame.
+ */
+private class ColumnPruner(columnsToPrune: Set[String]) extends
Transformer {
--- End diff --
Leave a TODO note to make this a public transformer.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
