[SPARK-54446][ML][CONNECT] FPGrowth supports local filesystem with Arrow file format

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

@@ -46,7 +46,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, SparkSession, SQLContext}
+import org.apache.spark.sql.execution.arrow.ArrowFileReadWrite
 import org.apache.spark.util.{Utils, VersionUtils}
 
 /**
@@ -1142,4 +1143,31 @@ 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)
+
+      df match {
+        case d: org.apache.spark.sql.classic.DataFrame =>
+          ArrowFileReadWrite.save(d, path)
+        case _ => throw new UnsupportedOperationException("Unsupported dataframe type")
+      }
+    } else {
+      df.write.parquet(path)
+    }
+  }
+
+  def loadDataFrame(path: String, spark: SparkSession): DataFrame = {
+    if (localSavingModeState.get()) {
+      spark match {
+        case s: org.apache.spark.sql.classic.SparkSession =>
+          ArrowFileReadWrite.load(s, path)
+        case _ => throw new UnsupportedOperationException("Unsupported session type")
+      }
+    } 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)
   }
 }
 

sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala

@@ -2340,14 +2340,13 @@ class Dataset[T] private[sql](
   }
 
   /** Convert to an RDD of serialized ArrowRecordBatches. */
-  private[sql] def toArrowBatchRdd(plan: SparkPlan): RDD[Array[Byte]] = {
+  private def toArrowBatchRddImpl(
+    plan: SparkPlan,
+    maxRecordsPerBatch: Int,
+    timeZoneId: String,
+    errorOnDuplicatedFieldNames: Boolean,
+    largeVarTypes: Boolean): RDD[Array[Byte]] = {
     val schemaCaptured = this.schema
-    val maxRecordsPerBatch = sparkSession.sessionState.conf.arrowMaxRecordsPerBatch
-    val timeZoneId = sparkSession.sessionState.conf.sessionLocalTimeZone
-    val errorOnDuplicatedFieldNames =
-      sparkSession.sessionState.conf.pandasStructHandlingMode == "legacy"
-    val largeVarTypes =
-      sparkSession.sessionState.conf.arrowUseLargeVarTypes
     plan.execute().mapPartitionsInternal { iter =>
       val context = TaskContext.get()
       ArrowConverters.toBatchIterator(
@@ -2361,7 +2360,24 @@ class Dataset[T] private[sql](
     }
   }
 
-  // This is only used in tests, for now.
+  private[sql] def toArrowBatchRdd(
+    maxRecordsPerBatch: Int,
+    timeZoneId: String,
+    errorOnDuplicatedFieldNames: Boolean,
+    largeVarTypes: Boolean): RDD[Array[Byte]] = {
+    toArrowBatchRddImpl(queryExecution.executedPlan,
+      maxRecordsPerBatch, timeZoneId, errorOnDuplicatedFieldNames, largeVarTypes)
+  }
+
+  private[sql] def toArrowBatchRdd(plan: SparkPlan): RDD[Array[Byte]] = {
+    toArrowBatchRddImpl(
+      plan,
+      sparkSession.sessionState.conf.arrowMaxRecordsPerBatch,
+      sparkSession.sessionState.conf.sessionLocalTimeZone,
+      sparkSession.sessionState.conf.pandasStructHandlingMode == "legacy",
+      sparkSession.sessionState.conf.arrowUseLargeVarTypes)
+  }
+
   private[sql] def toArrowBatchRdd: RDD[Array[Byte]] = {
     toArrowBatchRdd(queryExecution.executedPlan)
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala

@@ -555,25 +555,41 @@ private[sql] object ArrowConverters extends Logging {
       arrowBatches: Iterator[Array[Byte]],
       schemaString: String,
       session: SparkSession): DataFrame = {
-    val schema = DataType.fromJson(schemaString).asInstanceOf[StructType]
+    toDataFrame(
+      arrowBatches,
+      DataType.fromJson(schemaString).asInstanceOf[StructType],
+      session,
+      session.sessionState.conf.sessionLocalTimeZone,
+      false,
+      session.sessionState.conf.arrowUseLargeVarTypes)
+  }
+
+  /**
+   * Create a DataFrame from an iterator of serialized ArrowRecordBatches.
+   */
+  private[sql] def toDataFrame(
+      arrowBatches: Iterator[Array[Byte]],
+      schema: StructType,
+      session: SparkSession,
+      timeZoneId: String,
+      errorOnDuplicatedFieldNames: Boolean,
+      largeVarTypes: Boolean): DataFrame = {
     val attrs = toAttributes(schema)
     val batchesInDriver = arrowBatches.toArray
-    val largeVarTypes = session.sessionState.conf.arrowUseLargeVarTypes
     val shouldUseRDD = session.sessionState.conf
       .arrowLocalRelationThreshold < batchesInDriver.map(_.length.toLong).sum
 
     if (shouldUseRDD) {
       logDebug("Using RDD-based createDataFrame with Arrow optimization.")
-      val timezone = session.sessionState.conf.sessionLocalTimeZone
       val rdd = session.sparkContext
         .parallelize(batchesInDriver.toImmutableArraySeq, batchesInDriver.length)
         .mapPartitions { batchesInExecutors =>
           ArrowConverters.fromBatchIterator(
             batchesInExecutors,
             schema,
-            timezone,
-            errorOnDuplicatedFieldNames = false,
-            largeVarTypes = largeVarTypes,
+            timeZoneId,
+            errorOnDuplicatedFieldNames,
+            largeVarTypes,
             TaskContext.get())
         }
       session.internalCreateDataFrame(rdd.setName("arrow"), schema)
@@ -582,9 +598,9 @@ private[sql] object ArrowConverters extends Logging {
       val data = ArrowConverters.fromBatchIterator(
         batchesInDriver.iterator,
         schema,
-        session.sessionState.conf.sessionLocalTimeZone,
-        errorOnDuplicatedFieldNames = false,
-        largeVarTypes = largeVarTypes,
+        timeZoneId,
+        errorOnDuplicatedFieldNames,
+        largeVarTypes,
         TaskContext.get())
 
       // Project/copy it. Otherwise, the Arrow column vectors will be closed and released out.

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala

@@ -0,0 +1,110 @@
+/*
+ * 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.sql.execution.arrow
+
+import java.io.{ByteArrayOutputStream, FileOutputStream}
+import java.nio.channels.Channels
+import java.nio.file.Files
+import java.nio.file.Paths
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector._
+import org.apache.arrow.vector.ipc.{ArrowFileReader, ArrowFileWriter, WriteChannel}
+import org.apache.arrow.vector.ipc.message.MessageSerializer
+import org.apache.arrow.vector.types.pojo.Schema
+
+import org.apache.spark.sql.classic.{DataFrame, SparkSession}
+import org.apache.spark.sql.util.ArrowUtils
+
+private[sql] class SparkArrowFileWriter(
+  arrowSchema: Schema,
+  path: String) extends AutoCloseable {
+  private val allocator =
+    ArrowUtils.rootAllocator.newChildAllocator(
+      s"to${this.getClass.getSimpleName}", 0, Long.MaxValue)
+
+  protected val root = VectorSchemaRoot.create(arrowSchema, allocator)
+  protected val loader = new VectorLoader(root)
+  protected val arrowWriter = ArrowWriter.create(root)
+
+  protected val fileWriter =
+    new ArrowFileWriter(root, null, Channels.newChannel(new FileOutputStream(path)))
+
+  override def close(): Unit = {
+    root.close()
+    allocator.close()
+    fileWriter.close()
+  }
+
+  def write(batchBytesIter: Iterator[Array[Byte]]): Unit = {
+    fileWriter.start()
+    while (batchBytesIter.hasNext) {
+      val batchBytes = batchBytesIter.next()
+      val batch = ArrowConverters.loadBatch(batchBytes, allocator)
+      loader.load(batch)
+      fileWriter.writeBatch()
+    }
+    fileWriter.close()
+  }
+}
+
+private[sql] class SparkArrowFileReader(path: String) extends AutoCloseable {
+  private val allocator =
+    ArrowUtils.rootAllocator.newChildAllocator(
+      s"to${this.getClass.getSimpleName}", 0, Long.MaxValue)
+
+  protected val fileReader =
+    new ArrowFileReader(Files.newByteChannel(Paths.get(path)), allocator)
+
+  override def close(): Unit = {
+    allocator.close()
+    fileReader.close()
+  }
+
+  val schema: Schema = fileReader.getVectorSchemaRoot.getSchema
+
+  def read(): Iterator[Array[Byte]] = {
+    fileReader.getRecordBlocks.iterator().asScala.map { block =>
+      fileReader.loadRecordBatch(block)
+      val root = fileReader.getVectorSchemaRoot
+      val unloader = new VectorUnloader(root)
+      val batch = unloader.getRecordBatch
+      val out = new ByteArrayOutputStream()
+      val writeChannel = new WriteChannel(Channels.newChannel(out))
+      MessageSerializer.serialize(writeChannel, batch)
+      out.toByteArray
+    }
+  }
+}
+
+private[spark] object ArrowFileReadWrite {
+  def save(df: DataFrame, path: String): Unit = {
+    val maxRecordsPerBatch = df.sparkSession.sessionState.conf.arrowMaxRecordsPerBatch
+    val rdd = df.toArrowBatchRdd(maxRecordsPerBatch, "UTC", true, false)
+    val arrowSchema = ArrowUtils.toArrowSchema(df.schema, "UTC", true, false)
+    val writer = new SparkArrowFileWriter(arrowSchema, path)
+    writer.write(rdd.toLocalIterator)
+  }
+
+  def load(spark: SparkSession, path: String): DataFrame = {
+    val reader = new SparkArrowFileReader(path)
+    val schema = ArrowUtils.fromArrowSchema(reader.schema)
+    ArrowConverters.toDataFrame(reader.read(), schema, spark, "UTC", true, false)
+  }
+}