address comment

apache · Nov 15, 2018 · 5c4949d · 5c4949d
1 parent 2e3b089
commit 5c4949d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/fpm/PrefixSpan.scala b/mllib/src/main/scala/org/apache/spark/mllib/fpm/PrefixSpan.scala
@@ -174,7 +174,10 @@ class PrefixSpan private (
     val freqSequences = results.map { case (seq: Array[Int], count: Long) =>
       new FreqSequence(toPublicRepr(seq), count)
     }
+    // Cache the final RDD to the same storage level as input
+    freqSequences.persist(data.getStorageLevel)
     dataInternalRepr.unpersist(false)
+
     new PrefixSpanModel(freqSequences)
   }