fix

Author: zhengruifeng

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

@@ -343,16 +343,11 @@ object FPGrowthModel extends MLReadable[FPGrowthModel] {
   class FPGrowthModelWriter(instance: FPGrowthModel) extends MLWriter {
 
     override protected def saveImpl(path: String): Unit = {
-      if (ReadWriteUtils.localSavingModeState.get()) {
-        throw new UnsupportedOperationException(
-          "FPGrowthModel does not support saving to local filesystem path."
-        )
-      }
       val extraMetadata: JObject = Map("numTrainingRecords" -> instance.numTrainingRecords)
       DefaultParamsWriter.saveMetadata(instance, path, sparkSession,
         extraMetadata = Some(extraMetadata))
       val dataPath = new Path(path, "data").toString
-      instance.freqItemsets.write.parquet(dataPath)
+      ReadWriteUtils.saveDataFrame(dataPath, instance.freqItemsets)
     }
   }
 
@@ -362,11 +357,6 @@ object FPGrowthModel extends MLReadable[FPGrowthModel] {
     private val className = classOf[FPGrowthModel].getName
 
     override def load(path: String): FPGrowthModel = {
-      if (ReadWriteUtils.localSavingModeState.get()) {
-        throw new UnsupportedOperationException(
-          "FPGrowthModel does not support loading from local filesystem path."
-        )
-      }
       implicit val format = DefaultFormats
       val metadata = DefaultParamsReader.loadMetadata(path, sparkSession, className)
       val (major, minor) = VersionUtils.majorMinorVersion(metadata.sparkVersion)
@@ -378,7 +368,7 @@ object FPGrowthModel extends MLReadable[FPGrowthModel] {
         (metadata.metadata \ "numTrainingRecords").extract[Long]
       }
       val dataPath = new Path(path, "data").toString
-      val frequentItems = sparkSession.read.parquet(dataPath)
+      val frequentItems = ReadWriteUtils.loadDataFrame(dataPath, sparkSession)
       val itemSupport = if (numTrainingRecords == 0L) {
         Map.empty[Any, Double]
       } else {

mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala

@@ -19,10 +19,11 @@ package org.apache.spark.ml.util
 
 import java.io.{
   BufferedInputStream, BufferedOutputStream, DataInputStream, DataOutputStream,
-  File, FileInputStream, FileOutputStream, IOException
+  File, FileInputStream, FileOutputStream, IOException, ObjectInputStream,
+  ObjectOutputStream
 }
 import java.nio.file.{Files, Paths}
-import java.util.{Locale, ServiceLoader}
+import java.util.{ArrayList, Locale, ServiceLoader}
 
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
@@ -46,7 +47,8 @@ import org.apache.spark.ml.feature.RFormulaModel
 import org.apache.spark.ml.linalg.{DenseMatrix, DenseVector, Matrix, SparseMatrix, SparseVector, Vector}
 import org.apache.spark.ml.param.{ParamPair, Params}
 import org.apache.spark.ml.tuning.ValidatorParams
-import org.apache.spark.sql.{SparkSession, SQLContext}
+import org.apache.spark.sql.{DataFrame, Row, SparkSession, SQLContext}
+import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.{Utils, VersionUtils}
 
 /**
@@ -1142,4 +1144,47 @@ private[spark] object ReadWriteUtils {
       spark.read.parquet(path).as[T].collect()
     }
   }
+
+  def saveDataFrame(path: String, df: DataFrame): Unit = {
+    if (localSavingModeState.get()) {
+      val filePath = Paths.get(path)
+      Files.createDirectories(filePath.getParent)
+
+      Using.resource(
+        new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(filePath.toFile)))
+      ) { oos =>
+        val schema: StructType = df.schema
+        oos.writeObject(schema)
+        val it = df.toLocalIterator()
+        while (it.hasNext) {
+          oos.writeBoolean(true) // hasNext = True
+          val row: Row = it.next()
+          oos.writeObject(row)
+        }
+        oos.writeBoolean(false) // hasNext = False
+      }
+    } else {
+      df.write.parquet(path)
+    }
+  }
+
+  def loadDataFrame(path: String, spark: SparkSession): DataFrame = {
+    if (localSavingModeState.get()) {
+      Using.resource(
+        new ObjectInputStream(new BufferedInputStream(new FileInputStream(path)))
+      ) { ois =>
+        val schema = ois.readObject().asInstanceOf[StructType]
+        val rows = new ArrayList[Row]
+        var hasNext = ois.readBoolean()
+        while (hasNext) {
+          val row = ois.readObject().asInstanceOf[Row]
+          rows.add(row)
+          hasNext = ois.readBoolean()
+        }
+        spark.createDataFrame(rows, schema)
+      }
+    } else {
+      spark.read.parquet(path)
+    }
+  }
 }

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

@@ -165,7 +165,7 @@ class FPGrowthSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul
     }
     val fPGrowth = new FPGrowth()
     testEstimatorAndModelReadWrite(fPGrowth, dataset, FPGrowthSuite.allParamSettings,
-      FPGrowthSuite.allParamSettings, checkModelData, skipTestSaveLocal = true)
+      FPGrowthSuite.allParamSettings, checkModelData)
   }
 }