Github user ericl commented on a diff in the pull request:
https://github.com/apache/spark/pull/7574#discussion_r35598489
--- Diff: mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
---
@@ -62,19 +77,72 @@ class RFormula(override val uid: String)
/** @group getParam */
def getFormula: String = $(formula)
- /** @group getParam */
- def setFeaturesCol(value: String): this.type = set(featuresCol, value)
+ override def fit(dataset: DataFrame): RFormulaModel = {
+ require(parsedFormula.isDefined, "Must call setFormula() first.")
+ // StringType terms and terms representing interactions need to be
encoded before assembly.
+ // TODO(ekl) add support for feature interactions
+ var encoderStages = Seq[PipelineStage]()
+ var tempColumns = Seq[String]()
+ val encodedTerms = parsedFormula.get.terms.map { term =>
+ dataset.schema(term) match {
+ case column if column.dataType == StringType =>
+ val indexCol = term + "_idx_" + uid
+ val encodedCol = term + "_onehot_" + uid
+ encoderStages :+= new
StringIndexer().setInputCol(term).setOutputCol(indexCol)
+ encoderStages :+= new
OneHotEncoder().setInputCol(indexCol).setOutputCol(encodedCol)
+ tempColumns :+= indexCol
+ tempColumns :+= encodedCol
+ encodedCol
+ case _ =>
+ term
+ }
+ }
+ encoderStages :+= new VectorAssembler(uid)
+ .setInputCols(encodedTerms.toArray)
+ .setOutputCol($(featuresCol))
+ encoderStages :+= new ColumnPruner(tempColumns.toSet)
+ val pipelineModel = new
Pipeline(uid).setStages(encoderStages.toArray).fit(dataset)
+ copyValues(new RFormulaModel(uid, parsedFormula.get,
pipelineModel).setParent(this))
+ }
- /** @group getParam */
- def setLabelCol(value: String): this.type = set(labelCol, value)
+ // optimistic schema; does not contain any ML attributes
+ override def transformSchema(schema: StructType): StructType = {
+ if (hasLabelCol(schema)) {
+ StructType(schema.fields :+ StructField($(featuresCol), new
VectorUDT, true))
+ } else {
+ StructType(schema.fields :+ StructField($(featuresCol), new
VectorUDT, true) :+
+ StructField($(labelCol), DoubleType, true))
+ }
+ }
+
+ override def copy(extra: ParamMap): RFormula = defaultCopy(extra)
+
+ override def toString: String = s"RFormula(${get(formula)})"
+}
+
+/**
+ * A fitted RFormula. Fitting is required to determine the factor levels
of formula terms.
+ * @param parsedFormula a pre-parsed R formula.
+ * @param pipelineModel the fitted feature model, including factor to
index mappings.
+ */
+private[feature] class RFormulaModel(
--- End diff --
Done
---
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 [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
