[SPARK-8774] Add R model formula with basic support as a transformer

Author: ericl

mllib/src/main/scala/org/apache/spark/ml/feature/RModelFormula.scala

@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.feature
+
+import scala.util.parsing.combinator.RegexParsers
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.ml.Transformer
+import org.apache.spark.ml.param.{Param, ParamMap}
+import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasLabelCol}
+import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types._
+
+/**
+ * :: Experimental ::
+ * Implements the transforms required for fitting a dataset against a R model formula.
+ */
+@Experimental
+private[spark] class RModelFormula(override val uid: String)
+  extends Transformer with HasFeaturesCol with HasLabelCol {
+
+  def this() = this(Identifiable.randomUID("rModelFormula"))
+
+  val formula: Param[String] = new Param(this, "formula", "R model formula")
+  protected var parsedFormula: Option[RFormula] = None
+
+  def setFormula(value: String): this.type = {
+    parsedFormula = Some(RFormulaParser.parse(value))
+    set(formula, value)
+    this
+  }
+
+  override def transformSchema(schema: StructType): StructType = {
+    require(parsedFormula.isDefined, "Must call setFormula() first.")
+    val withFeatures = featureTransformer.transformSchema(schema)
+    val nullable = schema(parsedFormula.get.response).dataType match {
+      case _: NumericType | BooleanType => false
+      case _ => true
+    }
+    StructType(withFeatures.fields :+ StructField($(labelCol), DoubleType, nullable))
+  }
+ 
+  override def transform(dataset: DataFrame): DataFrame = {
+    require(parsedFormula.isDefined, "Must call setFormula() first.")
+    transformLabel(featureTransformer.transform(dataset))
+  }
+
+  override def copy(extra: ParamMap): RModelFormula = defaultCopy(extra)
+  
+  override def toString = s"RModelFormula(${get(formula)})"
+
+  protected def transformLabel(dataset: DataFrame): DataFrame = {
+    val responseName = parsedFormula.get.response
+    dataset.schema(responseName).dataType match {
+      case _: NumericType | BooleanType =>
+        dataset.select(
+          col("*"),
+          dataset(responseName).cast(DoubleType).as($(labelCol)))
+      case StringType =>
+        new StringIndexer(uid)
+          .setInputCol(responseName)
+          .setOutputCol($(labelCol))
+          .fit(dataset)
+          .transform(dataset)
+      case other =>
+        throw new IllegalArgumentException("Unsupported type for response: " + other)
+    }
+  }
+
+  protected def featureTransformer: Transformer = {
+    // TODO(ekl) add support for non-numeric features and feature interactions
+    new VectorAssembler(uid)
+      .setInputCols(parsedFormula.get.terms.toArray)
+      .setOutputCol($(featuresCol))
+  }
+}
+
+/**
+ * :: Experimental ::
+ * Represents a parsed R formula.
+ */
+private[ml] case class RFormula(response: String, terms: Seq[String])
+
+/**
+ * :: Experimental ::
+ * Limited implementation of R formula parsing. Currently supports: '~', '+'.
+ */
+private[ml] object RFormulaParser extends RegexParsers {
+  def term: Parser[String] = "([a-zA-Z]|\\.[a-zA-Z_])[a-zA-Z0-9._]*".r
+
+  def expr: Parser[List[String]] = term ~ rep("+" ~> term) ^^ { case a ~ list => a :: list }
+
+  def formula: Parser[RFormula] = (term ~ "~" ~ expr) ^^ { case r ~ "~" ~ t => RFormula(r, t) }
+
+  def parse(value: String): RFormula = parseAll(formula, value) match {
+    case Success(result, _) => result
+    case failure: NoSuccess => throw new IllegalArgumentException(
+      "Could not parse formula: " + value)
+  }
+}

mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala

@@ -116,7 +116,7 @@ class VectorAssembler(override val uid: String)
     if (schema.fieldNames.contains(outputColName)) {
       throw new IllegalArgumentException(s"Output column $outputColName already exists.")
     }
-    StructType(schema.fields :+ new StructField(outputColName, new VectorUDT, false))
+    StructType(schema.fields :+ new StructField(outputColName, new VectorUDT, true))
   }
 
   override def copy(extra: ParamMap): VectorAssembler = defaultCopy(extra)

mllib/src/test/scala/org/apache/spark/ml/feature/RModelFormulaSuite.scala

@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.feature
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.mllib.linalg.Vectors
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.mllib.util.TestingUtils._
+
+class RFormulaModelSuite extends SparkFunSuite with MLlibTestSparkContext {
+  test("params") {
+    ParamsSuite.checkParams(new RModelFormula())
+  }
+
+  test("parse simple formulas") {
+    def check(formula: String, response: String, terms: Seq[String]) {
+      new RModelFormula().setFormula(formula)
+      val parsed = RFormulaParser.parse(formula)
+      assert(parsed.response == response)
+      assert(parsed.terms == terms)
+    }
+    check("y ~ x", "y", Seq("x"))
+    check("y ~   ._foo  ", "y", Seq("._foo"))
+    check("resp ~ A_VAR + B + c123", "resp", Seq("A_VAR", "B", "c123"))
+  }
+
+  test("transform numeric data") {
+    val formula = new RModelFormula().setFormula("id ~ v1 + v2")
+    val original = sqlContext.createDataFrame(
+      Seq((0, 1.0, 3.0), (2, 2.0, 5.0))).toDF("id", "v1", "v2")
+    val result = formula.transform(original)
+    val resultSchema = formula.transformSchema(original.schema)
+    val expected = sqlContext.createDataFrame(
+      Seq(
+        (0, 1.0, 3.0, Vectors.dense(Array(1.0, 3.0)), 0.0),
+        (2, 2.0, 5.0, Vectors.dense(Array(2.0, 5.0)), 2.0))
+      ).toDF("id", "v1", "v2", "features", "label")
+    assert(result.schema.toString == resultSchema.toString)
+    assert(resultSchema.toString == expected.schema.toString)
+    assert(
+      result.collect.map(_.toString).mkString(",") ==
+      expected.collect.map(_.toString).mkString(","))
+  }
+
+  test("transform string label") {
+    val formula = new RModelFormula().setFormula("name ~ id")
+    val original = sqlContext.createDataFrame(
+      Seq((1, "foo"), (2, "bar"), (3, "bar"))).toDF("id", "name")
+    val result = formula.transform(original)
+    val resultSchema = formula.transformSchema(original.schema)
+    val expected = sqlContext.createDataFrame(
+      Seq(
+        (1, "foo", Vectors.dense(Array(1.0)), 1.0),
+        (2, "bar", Vectors.dense(Array(2.0)), 0.0),
+        (3, "bar", Vectors.dense(Array(3.0)), 0.0))
+      ).toDF("id", "name", "features", "label")
+    assert(result.schema.toString == resultSchema.toString)
+    assert(
+      result.collect.map(_.toString).mkString(",") ==
+      expected.collect.map(_.toString).mkString(","))
+  }
+}