[SPARK-14503][ML] spark.ml API for FPGrowth

## What changes were proposed in this pull request?

jira: https://issues.apache.org/jira/browse/SPARK-14503
Function parity: Add FPGrowth and AssociationRules to ML.

design doc: https://docs.google.com/document/d/1bVhABn5DiEj8bw0upqGMJT2L4nvO_0_cXdwu4uMT6uU/pub

Currently I make FPGrowthModel a transformer. For each association rule,  it will just examine the input items against antecedents and summarize the consequents.

Update:
Thinking again, FPGrowth is only the algorithm to find the frequent itemsets, and can be replaced by other algorithms. The frequent itemsets are used by AssociationRules to generate the association rules. Then we can use the association rules to predict with other records.

![drawing1](https://cloud.githubusercontent.com/assets/7981698/22489294/76b9302c-e7cb-11e6-8d2d-3fc53f407b2f.png)

**For reviewers**, Let's first decide if the current `transform` function meets your expectation.

Current options:

1. Current implementation: Use Estimator and Transformer pattern in ML, the `transform` function will examine the input items against all the association rules and summarize the consequents. Users can also access frequent items and association rules via other model members.

2. Keep the Estimator and Transformer pattern. But AssociationRulesModel and FPGrowthModel will have empty `transform` function, meaning DataFrame has no change after transform. But users can access frequent items and association rules via other model members.

3. (mentioned by zhengruifeng) Keep the Estimator and Transformer pattern. But `FPGrowthModel` and `AssociationRulesModel` will just return frequent itemsets and association rules DataFrame in the `transform` function. Meaning the resulting DataFrame after `transform` will not be related to the input DataFrame.

4. Discard the Estimator and Transformer pattern. Both FPGrowth and FPGrowthModel will directly extend from PipelineStage, thus we don't need to have a `transform` function.

 I'd like to hear more concrete suggestions. I would prefer option 1 or 2.

update 2:

As discussed  in the jira, we will not expose AssociationRules as a public API for now.

## How was this patch tested?

new unit test suites

Author: Yuhao <yuhao.yang@intel.com>
Author: Yuhao Yang <yuhao.yang@intel.com>
Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #15415 from hhbyyh/mlfpm.

mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala

+/*
+ * 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.fpm
+
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.ClassTag
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.param._
+import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasPredictionCol}
+import org.apache.spark.ml.util._
+import org.apache.spark.mllib.fpm.{AssociationRules => MLlibAssociationRules,
+  FPGrowth => MLlibFPGrowth}
+import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset
+import org.apache.spark.sql._
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types._
+
+/**
+ * Common params for FPGrowth and FPGrowthModel
+ */
+private[fpm] trait FPGrowthParams extends Params with HasFeaturesCol with HasPredictionCol {
+
+  /**
+   * Minimal support level of the frequent pattern. [0.0, 1.0]. Any pattern that appears
+   * more than (minSupport * size-of-the-dataset) times will be output
+   * Default: 0.3
+   * @group param
+   */
+  @Since("2.2.0")
+  val minSupport: DoubleParam = new DoubleParam(this, "minSupport",
+    "the minimal support level of a frequent pattern",
+    ParamValidators.inRange(0.0, 1.0))
+  setDefault(minSupport -> 0.3)
+
+  /** @group getParam */
+  @Since("2.2.0")
+  def getMinSupport: Double = $(minSupport)
+
+  /**
+   * Number of partitions (>=1) used by parallel FP-growth. By default the param is not set, and
+   * partition number of the input dataset is used.
+   * @group expertParam
+   */
+  @Since("2.2.0")
+  val numPartitions: IntParam = new IntParam(this, "numPartitions",
+    "Number of partitions used by parallel FP-growth", ParamValidators.gtEq[Int](1))
+
+  /** @group expertGetParam */
+  @Since("2.2.0")
+  def getNumPartitions: Int = $(numPartitions)
+
+  /**
+   * Minimal confidence for generating Association Rule.
+   * Note that minConfidence has no effect during fitting.
+   * Default: 0.8
+   * @group param
+   */
+  @Since("2.2.0")
+  val minConfidence: DoubleParam = new DoubleParam(this, "minConfidence",
+    "minimal confidence for generating Association Rule",
+    ParamValidators.inRange(0.0, 1.0))
+  setDefault(minConfidence -> 0.8)
+
+  /** @group getParam */
+  @Since("2.2.0")
+  def getMinConfidence: Double = $(minConfidence)
+
+  /**
+   * Validates and transforms the input schema.
+   * @param schema input schema
+   * @return output schema
+   */
+  @Since("2.2.0")
+  protected def validateAndTransformSchema(schema: StructType): StructType = {
+    val inputType = schema($(featuresCol)).dataType
+    require(inputType.isInstanceOf[ArrayType],
+      s"The input column must be ArrayType, but got $inputType.")
+    SchemaUtils.appendColumn(schema, $(predictionCol), schema($(featuresCol)).dataType)
+  }
+}
+
+/**
+ * :: Experimental ::
+ * A parallel FP-growth algorithm to mine frequent itemsets. The algorithm is described in
+ * <a href="http://dx.doi.org/10.1145/1454008.1454027">Li et al., PFP: Parallel FP-Growth for Query
+ * Recommendation</a>. PFP distributes computation in such a way that each worker executes an
+ * independent group of mining tasks. The FP-Growth algorithm is described in
+ * <a href="http://dx.doi.org/10.1145/335191.335372">Han et al., Mining frequent patterns without
+ * candidate generation</a>. Note null values in the feature column are ignored during fit().
+ *
+ * @see <a href="http://en.wikipedia.org/wiki/Association_rule_learning">
+ * Association rule learning (Wikipedia)</a>
+ */
+@Since("2.2.0")
+@Experimental
+class FPGrowth @Since("2.2.0") (
+    @Since("2.2.0") override val uid: String)
+  extends Estimator[FPGrowthModel] with FPGrowthParams with DefaultParamsWritable {
+
+  @Since("2.2.0")
+  def this() = this(Identifiable.randomUID("fpgrowth"))
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setMinSupport(value: Double): this.type = set(minSupport, value)
+
+  /** @group expertSetParam */
+  @Since("2.2.0")
+  def setNumPartitions(value: Int): this.type = set(numPartitions, value)
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setMinConfidence(value: Double): this.type = set(minConfidence, value)
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setFeaturesCol(value: String): this.type = set(featuresCol, value)
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setPredictionCol(value: String): this.type = set(predictionCol, value)
+
+  @Since("2.2.0")
+  override def fit(dataset: Dataset[_]): FPGrowthModel = {
+    transformSchema(dataset.schema, logging = true)
+    genericFit(dataset)
+  }
+
+  private def genericFit[T: ClassTag](dataset: Dataset[_]): FPGrowthModel = {
+    val data = dataset.select($(featuresCol))
+    val items = data.where(col($(featuresCol)).isNotNull).rdd.map(r => r.getSeq[T](0).toArray)
+    val mllibFP = new MLlibFPGrowth().setMinSupport($(minSupport))
+    if (isSet(numPartitions)) {
+      mllibFP.setNumPartitions($(numPartitions))
+    }
+    val parentModel = mllibFP.run(items)
+    val rows = parentModel.freqItemsets.map(f => Row(f.items, f.freq))
+
+    val schema = StructType(Seq(
+      StructField("items", dataset.schema($(featuresCol)).dataType, nullable = false),
+      StructField("freq", LongType, nullable = false)))
+    val frequentItems = dataset.sparkSession.createDataFrame(rows, schema)
+    copyValues(new FPGrowthModel(uid, frequentItems)).setParent(this)
+  }
+
+  @Since("2.2.0")
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+
+  @Since("2.2.0")
+  override def copy(extra: ParamMap): FPGrowth = defaultCopy(extra)
+}
+
+
+@Since("2.2.0")
+object FPGrowth extends DefaultParamsReadable[FPGrowth] {
+
+  @Since("2.2.0")
+  override def load(path: String): FPGrowth = super.load(path)
+}
+
+/**
+ * :: Experimental ::
+ * Model fitted by FPGrowth.
+ *
+ * @param freqItemsets frequent items in the format of DataFrame("items"[Seq], "freq"[Long])
+ */
+@Since("2.2.0")
+@Experimental
+class FPGrowthModel private[ml] (
+    @Since("2.2.0") override val uid: String,
+    @transient val freqItemsets: DataFrame)
+  extends Model[FPGrowthModel] with FPGrowthParams with MLWritable {
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setMinConfidence(value: Double): this.type = set(minConfidence, value)
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setFeaturesCol(value: String): this.type = set(featuresCol, value)
+
+  /** @group setParam */
+  @Since("2.2.0")
+  def setPredictionCol(value: String): this.type = set(predictionCol, value)
+
+  /**
+   * Get association rules fitted by AssociationRules using the minConfidence. Returns a dataframe
+   * with three fields, "antecedent", "consequent" and "confidence", where "antecedent" and
+   * "consequent" are Array[T] and "confidence" is Double.
+   */
+  @Since("2.2.0")
+  @transient lazy val associationRules: DataFrame = {
+    AssociationRules.getAssociationRulesFromFP(freqItemsets, "items", "freq", $(minConfidence))
+  }
+
+  /**
+   * The transform method first generates the association rules according to the frequent itemsets.
+   * Then for each association rule, it will examine the input items against antecedents and
+   * summarize the consequents as prediction. The prediction column has the same data type as the
+   * input column(Array[T]) and will not contain existing items in the input column. The null
+   * values in the feature columns are treated as empty sets.
+   * WARNING: internally it collects association rules to the driver and uses broadcast for
+   * efficiency. This may bring pressure to driver memory for large set of association rules.
+   */
+  @Since("2.2.0")
+  override def transform(dataset: Dataset[_]): DataFrame = {
+    transformSchema(dataset.schema, logging = true)
+    genericTransform(dataset)
+  }
+
+  private def genericTransform(dataset: Dataset[_]): DataFrame = {
+    val rules: Array[(Seq[Any], Seq[Any])] = associationRules.select("antecedent", "consequent")
+      .rdd.map(r => (r.getSeq(0), r.getSeq(1)))
+      .collect().asInstanceOf[Array[(Seq[Any], Seq[Any])]]
+    val brRules = dataset.sparkSession.sparkContext.broadcast(rules)
+
+    val dt = dataset.schema($(featuresCol)).dataType
+    // For each rule, examine the input items and summarize the consequents
+    val predictUDF = udf((items: Seq[_]) => {
+      if (items != null) {
+        val itemset = items.toSet
+        brRules.value.flatMap(rule =>
+          if (items != null && rule._1.forall(item => itemset.contains(item))) {
+            rule._2.filter(item => !itemset.contains(item))
+          } else {
+            Seq.empty
+          })
+      } else {
+        Seq.empty
+      }.distinct }, dt)
+    dataset.withColumn($(predictionCol), predictUDF(col($(featuresCol))))
+  }
+
+  @Since("2.2.0")
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+
+  @Since("2.2.0")
+  override def copy(extra: ParamMap): FPGrowthModel = {
+    val copied = new FPGrowthModel(uid, freqItemsets)
+    copyValues(copied, extra).setParent(this.parent)
+  }
+
+  @Since("2.2.0")
+  override def write: MLWriter = new FPGrowthModel.FPGrowthModelWriter(this)
+}
+
+@Since("2.2.0")
+object FPGrowthModel extends MLReadable[FPGrowthModel] {
+
+  @Since("2.2.0")
+  override def read: MLReader[FPGrowthModel] = new FPGrowthModelReader
+
+  @Since("2.2.0")
+  override def load(path: String): FPGrowthModel = super.load(path)
+
+  /** [[MLWriter]] instance for [[FPGrowthModel]] */
+  private[FPGrowthModel]
+  class FPGrowthModelWriter(instance: FPGrowthModel) extends MLWriter {
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val dataPath = new Path(path, "data").toString
+      instance.freqItemsets.write.parquet(dataPath)
+    }
+  }
+
+  private class FPGrowthModelReader extends MLReader[FPGrowthModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[FPGrowthModel].getName
+
+    override def load(path: String): FPGrowthModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val frequentItems = sparkSession.read.parquet(dataPath)
+      val model = new FPGrowthModel(metadata.uid, frequentItems)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+}
+
+private[fpm] object AssociationRules {
+
+  /**
+   * Computes the association rules with confidence above minConfidence.
+   * @param dataset DataFrame("items", "freq") containing frequent itemset obtained from
+   *                algorithms like [[FPGrowth]].
+   * @param itemsCol column name for frequent itemsets
+   * @param freqCol column name for frequent itemsets count
+   * @param minConfidence minimum confidence for the result association rules
+   * @return a DataFrame("antecedent", "consequent", "confidence") containing the association
+   *         rules.
+   */
+  def getAssociationRulesFromFP[T: ClassTag](
+        dataset: Dataset[_],
+        itemsCol: String,
+        freqCol: String,
+        minConfidence: Double): DataFrame = {
+
+    val freqItemSetRdd = dataset.select(itemsCol, freqCol).rdd
+      .map(row => new FreqItemset(row.getSeq[T](0).toArray, row.getLong(1)))
+    val rows = new MLlibAssociationRules()
+      .setMinConfidence(minConfidence)
+      .run(freqItemSetRdd)
+      .map(r => Row(r.antecedent, r.consequent, r.confidence))
+
+    val dt = dataset.schema(itemsCol).dataType
+    val schema = StructType(Seq(
+      StructField("antecedent", dt, nullable = false),
+      StructField("consequent", dt, nullable = false),
+      StructField("confidence", DoubleType, nullable = false)))
+    val rules = dataset.sparkSession.createDataFrame(rows, schema)
+    rules
+  }
+}

mllib/src/test/scala/org/apache/spark/ml/fpm/FPGrowthSuite.scala

+/*
+ * 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.fpm
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types._
+
+class FPGrowthSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
+  @transient var dataset: Dataset[_] = _
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    dataset = FPGrowthSuite.getFPGrowthData(spark)
+  }
+
+  test("FPGrowth fit and transform with different data types") {
+    Array(IntegerType, StringType, ShortType, LongType, ByteType).foreach { dt =>
+      val data = dataset.withColumn("features", col("features").cast(ArrayType(dt)))
+      val model = new FPGrowth().setMinSupport(0.5).fit(data)
+      val generatedRules = model.setMinConfidence(0.5).associationRules
+      val expectedRules = spark.createDataFrame(Seq(
+        (Array("2"), Array("1"), 1.0),
+        (Array("1"), Array("2"), 0.75)
+      )).toDF("antecedent", "consequent", "confidence")
+        .withColumn("antecedent", col("antecedent").cast(ArrayType(dt)))
+        .withColumn("consequent", col("consequent").cast(ArrayType(dt)))
+      assert(expectedRules.sort("antecedent").rdd.collect().sameElements(
+        generatedRules.sort("antecedent").rdd.collect()))
+
+      val transformed = model.transform(data)
+      val expectedTransformed = spark.createDataFrame(Seq(
+        (0, Array("1", "2"), Array.emptyIntArray),
+        (0, Array("1", "2"), Array.emptyIntArray),
+        (0, Array("1", "2"), Array.emptyIntArray),
+        (0, Array("1", "3"), Array(2))
+      )).toDF("id", "features", "prediction")
+        .withColumn("features", col("features").cast(ArrayType(dt)))
+        .withColumn("prediction", col("prediction").cast(ArrayType(dt)))
+      assert(expectedTransformed.collect().toSet.equals(
+        transformed.collect().toSet))
+    }
+  }
+
+  test("FPGrowth getFreqItems") {
+    val model = new FPGrowth().setMinSupport(0.7).fit(dataset)
+    val expectedFreq = spark.createDataFrame(Seq(
+      (Array("1"), 4L),
+      (Array("2"), 3L),
+      (Array("1", "2"), 3L),
+      (Array("2", "1"), 3L) // duplicate as the items sequence is not guaranteed
+    )).toDF("items", "expectedFreq")
+    val freqItems = model.freqItemsets
+
+    val checkDF = freqItems.join(expectedFreq, "items")
+    assert(checkDF.count() == 3 && checkDF.filter(col("freq") === col("expectedFreq")).count() == 3)
+  }
+
+  test("FPGrowth getFreqItems with Null") {
+    val df = spark.createDataFrame(Seq(
+      (1, Array("1", "2", "3", "5")),
+      (2, Array("1", "2", "3", "4")),
+      (3, null.asInstanceOf[Array[String]])
+    )).toDF("id", "features")
+    val model = new FPGrowth().setMinSupport(0.7).fit(dataset)
+    val prediction = model.transform(df)
+    assert(prediction.select("prediction").where("id=3").first().getSeq[String](0).isEmpty)
+  }
+
+  test("FPGrowth parameter check") {
+    val fpGrowth = new FPGrowth().setMinSupport(0.4567)
+    val model = fpGrowth.fit(dataset)
+      .setMinConfidence(0.5678)
+    assert(fpGrowth.getMinSupport === 0.4567)
+    assert(model.getMinConfidence === 0.5678)
+  }
+
+  test("read/write") {
+    def checkModelData(model: FPGrowthModel, model2: FPGrowthModel): Unit = {
+      assert(model.freqItemsets.sort("items").collect() ===
+        model2.freqItemsets.sort("items").collect())
+    }
+    val fPGrowth = new FPGrowth()
+    testEstimatorAndModelReadWrite(
+      fPGrowth, dataset, FPGrowthSuite.allParamSettings, checkModelData)
+  }
+
+}
+
+object FPGrowthSuite {
+
+  def getFPGrowthData(spark: SparkSession): DataFrame = {
+    spark.createDataFrame(Seq(
+      (0, Array("1", "2")),
+      (0, Array("1", "2")),
+      (0, Array("1", "2")),
+      (0, Array("1", "3"))
+    )).toDF("id", "features")
+  }
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "minSupport" -> 0.321,
+    "minConfidence" -> 0.456,
+    "numPartitions" -> 5,
+    "predictionCol" -> "myPrediction"
+  )
+}