apply arrow

nit

Author: zhengruifeng

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

@@ -17,7 +17,7 @@
 
 package org.apache.spark.ml.util
 
-import org.apache.spark.SparkException
+import org.apache.spark.{SparkException, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.LogKeys.{CLASS_NAME, LABEL_COLUMN, NUM_CLASSES}
 import org.apache.spark.ml.PredictorParams
@@ -28,6 +28,7 @@ import org.apache.spark.ml.param.shared.HasWeightCol
 import org.apache.spark.mllib.linalg.{Vector => OldVector, Vectors => OldVectors}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
+import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
 
@@ -212,4 +213,51 @@ private[spark] object DatasetUtils extends Logging {
       dataset.select(columnToVector(dataset, vectorCol)).head().getAs[Vector](0).size
     }
   }
+
+  private[ml] def toArrowBatchRDD(
+      dataFrame: DataFrame,
+      timeZoneId: String): RDD[Array[Byte]] = {
+    dataFrame match {
+      case df: org.apache.spark.sql.classic.DataFrame =>
+        val spark = df.sparkSession
+        val schema = df.schema
+        val maxRecordsPerBatch = spark.sessionState.conf.arrowMaxRecordsPerBatch
+        df.queryExecution.executedPlan.execute().mapPartitionsInternal { iter =>
+          val context = TaskContext.get()
+          ArrowConverters.toBatchIterator(
+            iter,
+            schema,
+            maxRecordsPerBatch,
+            timeZoneId,
+            true,
+            false,
+            context)
+        }
+
+      case _ => throw new UnsupportedOperationException("Not implemented")
+    }
+  }
+
+  private[ml] def fromArrowBatchRDD(
+      rdd: RDD[Array[Byte]],
+      schema: StructType,
+      timeZoneId: String,
+      sparkSession: SparkSession): DataFrame = {
+    sparkSession match {
+      case spark: org.apache.spark.sql.classic.SparkSession =>
+        val rowRDD = rdd.mapPartitions { iter =>
+          val context = TaskContext.get()
+          ArrowConverters.fromBatchIterator(
+            iter,
+            schema,
+            timeZoneId,
+            true,
+            false,
+            context)
+        }
+        spark.internalCreateDataFrame(rowRDD.setName("arrow"), schema)
+
+      case _ => throw new UnsupportedOperationException("Not implemented")
+    }
+  }
 }

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

@@ -19,11 +19,10 @@ package org.apache.spark.ml.util
 
 import java.io.{
   BufferedInputStream, BufferedOutputStream, DataInputStream, DataOutputStream,
-  File, FileInputStream, FileOutputStream, IOException, ObjectInputStream,
-  ObjectOutputStream
+  File, FileInputStream, FileOutputStream, IOException
 }
 import java.nio.file.{Files, Paths}
-import java.util.{ArrayList, Locale, ServiceLoader}
+import java.util.{Locale, ServiceLoader}
 
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
@@ -47,7 +46,7 @@ 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.{DataFrame, Row, SparkSession, SQLContext}
+import org.apache.spark.sql.{DataFrame, SparkSession, SQLContext}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.{Utils, VersionUtils}
 
@@ -1151,17 +1150,21 @@ private[spark] object ReadWriteUtils {
       Files.createDirectories(filePath.getParent)
 
       Using.resource(
-        new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(filePath.toFile)))
-      ) { oos =>
+        new DataOutputStream(new BufferedOutputStream(new FileOutputStream(filePath.toFile)))
+      ) { dos =>
+        dos.writeUTF("ARROW") // format
+
         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)
+        dos.writeUTF(schema.json)
+
+        val iter = DatasetUtils.toArrowBatchRDD(df, "UTC").toLocalIterator
+        while (iter.hasNext) {
+          val bytes = iter.next()
+          require(bytes != null)
+          dos.writeInt(bytes.length)
+          dos.write(bytes)
         }
-        oos.writeBoolean(false) // hasNext = False
+        dos.writeInt(-1) // End
       }
     } else {
       df.write.parquet(path)
@@ -1170,18 +1173,33 @@ private[spark] object ReadWriteUtils {
 
   def loadDataFrame(path: String, spark: SparkSession): DataFrame = {
     if (localSavingModeState.get()) {
+      val sc = spark match {
+        case s: org.apache.spark.sql.classic.SparkSession => s.sparkContext
+      }
+
       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()
+        new DataInputStream(new BufferedInputStream(new FileInputStream(path)))
+      ) { dis =>
+        val format = dis.readUTF()
+        require(format == "ARROW")
+
+        val schema: StructType = StructType.fromString(dis.readUTF())
+
+        val buff = mutable.ListBuffer.empty[Array[Byte]]
+        var nextBytes = dis.readInt()
+        while (nextBytes >= 0) {
+          val bytes = dis.readNBytes(nextBytes)
+          buff.append(bytes)
+          nextBytes = dis.readInt()
         }
-        spark.createDataFrame(rows, schema)
+        require(nextBytes == -1)
+
+        DatasetUtils.fromArrowBatchRDD(
+          sc.parallelize[Array[Byte]](buff.result()),
+          schema,
+          "UTC",
+          spark
+        )
       }
     } else {
       spark.read.parquet(path)

sql/api/src/main/scala/org/apache/spark/sql/types/StructType.scala

@@ -531,7 +531,7 @@ object StructType extends AbstractDataType {
 
   override private[sql] def simpleString: String = "struct"
 
-  private[sql] def fromString(raw: String): StructType = {
+  private[spark] def fromString(raw: String): StructType = {
     Try(DataType.fromJson(raw)).getOrElse(LegacyTypeStringParser.parseString(raw)) match {
       case t: StructType => t
       case _ => throw DataTypeErrors.failedParsingStructTypeError(raw)

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

@@ -399,7 +399,7 @@ class SparkSession private(
   /**
    * Creates a `DataFrame` from an `RDD[InternalRow]`.
    */
-  private[sql] def internalCreateDataFrame(
+  private[spark] def internalCreateDataFrame(
       catalystRows: RDD[InternalRow],
       schema: StructType,
       isStreaming: Boolean = false): DataFrame = {

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

@@ -80,7 +80,7 @@ private[sql] class ArrowBatchStreamWriter(
   }
 }
 
-private[sql] object ArrowConverters extends Logging {
+private[spark] object ArrowConverters extends Logging {
   private[sql] class ArrowBatchIterator(
       rowIter: Iterator[InternalRow],
       schema: StructType,
@@ -231,7 +231,7 @@ private[sql] object ArrowConverters extends Logging {
    * Maps Iterator from InternalRow to serialized ArrowRecordBatches. Limit ArrowRecordBatch size
    * in a batch by setting maxRecordsPerBatch or use 0 to fully consume rowIter.
    */
-  private[sql] def toBatchIterator(
+  private[spark] def toBatchIterator(
       rowIter: Iterator[InternalRow],
       schema: StructType,
       maxRecordsPerBatch: Long,
@@ -484,7 +484,7 @@ private[sql] object ArrowConverters extends Logging {
   /**
    * Maps iterator from serialized ArrowRecordBatches to InternalRows.
    */
-  private[sql] def fromBatchIterator(
+  private[spark] def fromBatchIterator(
       arrowBatchIter: Iterator[Array[Byte]],
       schema: StructType,
       timeZoneId: String,