add support to association rules

Author: YY-OnCall

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

@@ -17,10 +17,11 @@
 
 package org.apache.spark.ml.fpm
 
-import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
 
 import org.apache.hadoop.fs.Path
+import org.json4s.DefaultFormats
+import org.json4s.JsonDSL._
 
 import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{Estimator, Model}
@@ -190,7 +191,7 @@ object FPGrowth extends DefaultParamsReadable[FPGrowth] {
 class FPGrowthModel private[ml] (
     @Since("2.2.0") override val uid: String,
     @transient val freqItemsets: DataFrame,
-    val numTotalRecords: Long)
+    val numTrainingRecords: Long)
   extends Model[FPGrowthModel] with FPGrowthParams with MLWritable {
 
   /** @group setParam */
@@ -212,7 +213,8 @@ class FPGrowthModel private[ml] (
    */
   @Since("2.2.0")
   @transient lazy val associationRules: DataFrame = {
-    AssociationRules.getAssociationRulesFromFP(freqItemsets, "items", "freq", numTotalRecords, $(minConfidence))
+    AssociationRules.getAssociationRulesFromFP(
+      freqItemsets, "items", "freq", numTrainingRecords, $(minConfidence))
   }
 
   /**
@@ -260,7 +262,7 @@ class FPGrowthModel private[ml] (
 
   @Since("2.2.0")
   override def copy(extra: ParamMap): FPGrowthModel = {
-    val copied = new FPGrowthModel(uid, freqItemsets)
+    val copied = new FPGrowthModel(uid, freqItemsets, numTrainingRecords)
     copyValues(copied, extra).setParent(this.parent)
   }
 
@@ -282,7 +284,8 @@ object FPGrowthModel extends MLReadable[FPGrowthModel] {
   class FPGrowthModelWriter(instance: FPGrowthModel) extends MLWriter {
 
     override protected def saveImpl(path: String): Unit = {
-      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val extraMetadata = "numTrainingRecords" -> instance.numTrainingRecords
+      DefaultParamsWriter.saveMetadata(instance, path, sc, Some(extraMetadata))
       val dataPath = new Path(path, "data").toString
       instance.freqItemsets.write.parquet(dataPath)
     }
@@ -295,9 +298,11 @@ object FPGrowthModel extends MLReadable[FPGrowthModel] {
 
     override def load(path: String): FPGrowthModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      implicit val format = DefaultFormats
+      val numTrainingRecords = (metadata.metadata \ "numTrainingRecords").extract[Long]
       val dataPath = new Path(path, "data").toString
       val frequentItems = sparkSession.read.parquet(dataPath)
-      val model = new FPGrowthModel(metadata.uid, frequentItems)
+      val model = new FPGrowthModel(metadata.uid, frequentItems, numTrainingRecords)
       DefaultParamsReader.getAndSetParams(model, metadata)
       model
     }
@@ -312,25 +317,24 @@ private[fpm] object AssociationRules {
    *                algorithms like [[FPGrowth]].
    * @param itemsCol column name for frequent itemsets
    * @param freqCol column name for frequent itemsets count
+   * @param numTrainingRecords count of training Dataset
    * @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,
-        numTotalRecords: Long,
-        minConfidence: Double): DataFrame = {
+      dataset: Dataset[_],
+      itemsCol: String,
+      freqCol: String,
+      numTrainingRecords: Long,
+      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, r.freqUnion / numTotalRecords))
-
-
+      .map(r => Row(r.antecedent, r.consequent, r.confidence, r.freqUnion / numTrainingRecords))
 
     val dt = dataset.schema(itemsCol).dataType
     val schema = StructType(Seq(

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

@@ -38,13 +38,13 @@ class FPGrowthSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul
       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")
+        (Array("2"), Array("1"), 1.0, 0.75),
+        (Array("1"), Array("2"), 0.75, 0.75)
+      )).toDF("antecedent", "consequent", "confidence", "support")
         .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()))
+      assert(expectedRules.collect().toSet.equals(
+        generatedRules.collect().toSet))
 
       val transformed = model.transform(data)
       val expectedTransformed = spark.createDataFrame(Seq(
@@ -74,6 +74,20 @@ class FPGrowthSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul
     assert(checkDF.count() == 3 && checkDF.filter(col("freq") === col("expectedFreq")).count() == 3)
   }
 
+  test("FPGrowth associationRules") {
+    val freqItemsets = spark.createDataFrame(Seq(
+      (Array("2"), 4L),
+      (Array("1"), 2L),
+      (Array("1", "2"), 2L)
+    )).toDF("items", "freq")
+    val model = new FPGrowthModel("fpgrowth", freqItemsets, 4L).setMinConfidence(0.1)
+    val expectedRules = spark.createDataFrame(Seq(
+      (Array("2"), Array("1"), 0.5, 0.5),
+      (Array("1"), Array("2"), 1.0, 0.5)
+    )).toDF("antecedent", "consequent", "confidence", "support")
+    assert(expectedRules.collect().toSet.equals(model.associationRules.collect().toSet))
+  }
+
   test("FPGrowth getFreqItems with Null") {
     val df = spark.createDataFrame(Seq(
       (1, Array("1", "2", "3", "5")),