[SPARK-8998][MLlib] Collect enough frequent prefixes before local processing in PrefixSpan (new)

mllib/src/main/scala/org/apache/spark/mllib/fpm/PrefixSpan.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.mllib.fpm
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.apache.spark.Logging
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
@@ -43,6 +45,8 @@ class PrefixSpan private (
     private var minSupport: Double,
     private var maxPatternLength: Int) extends Logging with Serializable {
 
+  private val maxSuffixesBeforeLocalProcessing: Long = 10000
+
   /**
    * Constructs a default instance with default parameters
    * {minSupport: `0.1`, maxPatternLength: `10`}.
@@ -82,20 +86,106 @@ class PrefixSpan private (
       logWarning("Input data is not cached.")
     }
     val minCount = getMinCount(sequences)
-    val lengthOnePatternsAndCounts =
-      getFreqItemAndCounts(minCount, sequences).collect()
-    val prefixAndProjectedDatabase = getPrefixAndProjectedDatabase(
-      lengthOnePatternsAndCounts.map(_._1), sequences)
-    val groupedProjectedDatabase = prefixAndProjectedDatabase
-      .map(x => (x._1.toSeq, x._2))
+    val lengthOnePatternsAndCounts = getFreqItemAndCounts(minCount, sequences)
+    val prefixSuffixPairs = getPrefixSuffixPairs(
+      lengthOnePatternsAndCounts.map(_._1).collect(), sequences)
+    var patternsCount: Long = lengthOnePatternsAndCounts.count()
+    var allPatternAndCounts = lengthOnePatternsAndCounts.map(x => (ArrayBuffer(x._1), x._2))
+    var (smallPrefixSuffixPairs, largePrefixSuffixPairs) =
+      splitPrefixSuffixPairs(prefixSuffixPairs)
+    largePrefixSuffixPairs.persist(StorageLevel.MEMORY_AND_DISK)
+    var patternLength: Int = 1
+    while (patternLength < maxPatternLength &&
+      largePrefixSuffixPairs.count() != 0) {
+      val (nextPatternAndCounts, nextPrefixSuffixPairs) =
+        getPatternCountsAndPrefixSuffixPairs(minCount, largePrefixSuffixPairs)
+      patternsCount = nextPatternAndCounts.count()
+      largePrefixSuffixPairs.unpersist()
+      val splitedPrefixSuffixPairs = splitPrefixSuffixPairs(nextPrefixSuffixPairs)
+      largePrefixSuffixPairs = splitedPrefixSuffixPairs._2
+      largePrefixSuffixPairs.persist(StorageLevel.MEMORY_AND_DISK)
+      smallPrefixSuffixPairs = smallPrefixSuffixPairs ++ splitedPrefixSuffixPairs._1
+      allPatternAndCounts = allPatternAndCounts ++ nextPatternAndCounts
+      patternLength = patternLength + 1
+    }
+    if (smallPrefixSuffixPairs.count() > 0) {
+      val projectedDatabase = smallPrefixSuffixPairs
+        .map(x => (x._1.toSeq, x._2))
+        .groupByKey()
+        .map(x => (x._1.toArray, x._2.toArray))
+      val nextPatternAndCounts = getPatternsInLocal(minCount, projectedDatabase)
+      allPatternAndCounts = allPatternAndCounts ++ nextPatternAndCounts
+    }
+    allPatternAndCounts.map { case (pattern, count) => (pattern.toArray, count) }
+  }
+
+
+  /**
+   * Split prefix suffix pairs to two parts:
+   * suffixes' size less than maxSuffixesBeforeLocalProcessing and
+   * suffixes' size more than maxSuffixesBeforeLocalProcessing
+   * @param prefixSuffixPairs prefix (length n) and suffix pairs,
+   * @return small size prefix suffix pairs and big size prefix suffix pairs
+   *         (RDD[prefix, suffix], RDD[prefix, suffix ])
+   */
+  private def splitPrefixSuffixPairs(
+      prefixSuffixPairs: RDD[(ArrayBuffer[Int], Array[Int])]):
+  (RDD[(ArrayBuffer[Int], Array[Int])], RDD[(ArrayBuffer[Int], Array[Int])]) = {
+    val suffixSizeMap = prefixSuffixPairs
+      .map(x => (x._1, x._2.length))
+      .reduceByKey(_ + _)
+      .map(x => (x._2 <= maxSuffixesBeforeLocalProcessing, Set(x._1)))
+      .reduceByKey(_ ++ _)
+      .collect
+      .toMap
+    val small = if (suffixSizeMap.contains(true)) {
+      prefixSuffixPairs.filter(x => suffixSizeMap(true).contains(x._1))
+    } else {
+      prefixSuffixPairs.filter(x => false)
+    }
+    val large = if (suffixSizeMap.contains(false)) {
+      prefixSuffixPairs.filter(x => suffixSizeMap(false).contains(x._1))
+    } else {
+      prefixSuffixPairs.filter(x => false)
+    }
+    (small, large)
+  }
+
+  /**
+   * Get the pattern and counts, and prefix suffix pairs
+   * @param minCount minimum count
+   * @param prefixSuffixPairs prefix (length n) and suffix pairs,
+   * @return pattern (length n+1) and counts, and prefix (length n+1) and suffix pairs
+   *         (RDD[pattern, count], RDD[prefix, suffix ])
+   */
+  private def getPatternCountsAndPrefixSuffixPairs(
+      minCount: Long,
+      prefixSuffixPairs: RDD[(ArrayBuffer[Int], Array[Int])]):
+  (RDD[(ArrayBuffer[Int], Long)], RDD[(ArrayBuffer[Int], Array[Int])]) = {
+    val prefixAndFrequentItemAndCounts = prefixSuffixPairs
+      .flatMap { case (prefix, suffix) => suffix.distinct.map(y => ((prefix, y), 1L)) }
+      .reduceByKey(_ + _)
+      .filter(_._2 >= minCount)
+    val patternAndCounts = prefixAndFrequentItemAndCounts
+      .map { case ((prefix, item), count) => (prefix :+ item, count) }
+    val prefixToFrequentNextItemsMap = prefixAndFrequentItemAndCounts
+      .keys
       .groupByKey()
-      .map(x => (x._1.toArray, x._2.toArray))
-    val nextPatterns = getPatternsInLocal(minCount, groupedProjectedDatabase)
-    val lengthOnePatternsAndCountsRdd =
-      sequences.sparkContext.parallelize(
-        lengthOnePatternsAndCounts.map(x => (Array(x._1), x._2)))
-    val allPatterns = lengthOnePatternsAndCountsRdd ++ nextPatterns
-    allPatterns
+      .mapValues(_.toSet)
+      .collect()
+      .toMap
+    val nextPrefixSuffixPairs = prefixSuffixPairs
+      .filter(x => prefixToFrequentNextItemsMap.contains(x._1))
+      .flatMap { case (prefix, suffix) =>
+      val frequentNextItems = prefixToFrequentNextItemsMap(prefix)
+      val filteredSuffix = suffix.filter(frequentNextItems.contains(_))
+      frequentNextItems.flatMap { item =>
+        val suffix = LocalPrefixSpan.getSuffix(item, filteredSuffix)
+        if (suffix.isEmpty) None
+        else Some(prefix :+ item, suffix)
+      }
+    }
+    (patternAndCounts, nextPrefixSuffixPairs)
   }
 
   /**
@@ -122,37 +212,40 @@ class PrefixSpan private (
   }
 
   /**
-   * Get the frequent prefixes' projected database.
+   * Get the frequent prefixes and suffix pairs.
    * @param frequentPrefixes frequent prefixes
    * @param sequences sequences data
-   * @return prefixes and projected database
+   * @return prefixes and suffix pairs.
    */
-  private def getPrefixAndProjectedDatabase(
+  private def getPrefixSuffixPairs(
       frequentPrefixes: Array[Int],
-      sequences: RDD[Array[Int]]): RDD[(Array[Int], Array[Int])] = {
+      sequences: RDD[Array[Int]]): RDD[(ArrayBuffer[Int], Array[Int])] = {
     val filteredSequences = sequences.map { p =>
       p.filter (frequentPrefixes.contains(_) )
     }
     filteredSequences.flatMap { x =>
       frequentPrefixes.map { y =>
         val sub = LocalPrefixSpan.getSuffix(y, x)
-        (Array(y), sub)
+        (ArrayBuffer(y), sub)
       }.filter(_._2.nonEmpty)
     }
   }
 
   /**
    * calculate the patterns in local.
    * @param minCount the absolute minimum count
-   * @param data patterns and projected sequences data data
+   * @param data prefixes and projected sequences data data
    * @return patterns
    */
   private def getPatternsInLocal(
       minCount: Long,
-      data: RDD[(Array[Int], Array[Array[Int]])]): RDD[(Array[Int], Long)] = {
-    data.flatMap { case (prefix, projDB) =>
-      LocalPrefixSpan.run(minCount, maxPatternLength, prefix.toList, projDB)
-        .map { case (pattern: List[Int], count: Long) => (pattern.toArray.reverse, count) }
+      data: RDD[(Array[Int], Array[Array[Int]])]): RDD[(ArrayBuffer[Int], Long)] = {
+    data.flatMap {
+    case (prefix, projDB) =>
+      LocalPrefixSpan.run(minCount, maxPatternLength, prefix.toList.reverse, projDB)
+        .map { case (pattern: List[Int], count: Long) =>
+        (pattern.toArray.reverse.to[ArrayBuffer], count)
+      }
     }
   }
 }