arrow to rdd

pang-wu · pang-wu · commit d27c2d40802f · 2025-08-26T12:47:04.000-07:00
diff --git a/core/raydp-main/src/main/scala/org/apache/spark/sql/raydp/ObjectStoreWriter.scala b/core/raydp-main/src/main/scala/org/apache/spark/sql/raydp/ObjectStoreWriter.scala
@@ -225,7 +225,7 @@ object ObjectStoreWriter {
     }
     val uuid = dfToId.getOrElseUpdate(df, UUID.randomUUID())
     val queue = ObjectRefHolder.getQueue(uuid)
-    val rdd = df.toArrowBatchRdd
+    val rdd = SparkShimLoader.getSparkShims.toArrowBatchRDD(df)
     rdd.persist(storageLevel)
     rdd.count()
     var executorIds = df.sqlContext.sparkContext.getExecutorIds.toArray
diff --git a/core/shims/common/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/common/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -21,6 +21,7 @@ import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.executor.RayDPExecutorBackendFactory
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SparkSession}
 
@@ -40,4 +41,6 @@ trait SparkShims {
   def getDummyTaskContext(partitionId: Int, env: SparkEnv): TaskContext
 
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema
+
+  def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]]
 }
diff --git a/core/shims/spark322/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/spark322/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.spark322.SparkSqlUtils
 import com.intel.raydp.shims.{ShimDescriptor, SparkShims}
 import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.StructType
 
 class Spark322Shims extends SparkShims {
@@ -49,4 +50,8 @@ class Spark322Shims extends SparkShims {
   override def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     SparkSqlUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  override def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]] = {
+    SparkSqlUtils.toArrowRDD(dataFrame, dataFrame.sparkSession)
+  }
 }
diff --git a/core/shims/spark322/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala b/core/shims/spark322/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.spark322
 
 import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, SQLContext, SparkSession}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types.StructType
@@ -32,4 +33,8 @@ object SparkSqlUtils {
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     ArrowUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  def toArrowRDD(dataFrame: DataFrame, sparkSession: SparkSession): RDD[Array[Byte]] = {
+    dataFrame.toArrowBatchRdd
+  }
 }
diff --git a/core/shims/spark330/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/spark330/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.spark330.SparkSqlUtils
 import com.intel.raydp.shims.{ShimDescriptor, SparkShims}
 import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.StructType
 
 class Spark330Shims extends SparkShims {
@@ -49,4 +50,8 @@ class Spark330Shims extends SparkShims {
   override def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     SparkSqlUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  override def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]] = {
+    SparkSqlUtils.toArrowRDD(dataFrame, dataFrame.sparkSession)
+  }
 }
diff --git a/core/shims/spark330/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala b/core/shims/spark330/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.spark330
 
 import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, SQLContext, SparkSession}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types.StructType
@@ -32,4 +33,8 @@ object SparkSqlUtils {
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     ArrowUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  def toArrowRDD(dataFrame: DataFrame, sparkSession: SparkSession): RDD[Array[Byte]] = {
+    dataFrame.toArrowBatchRdd
+  }
 }
diff --git a/core/shims/spark340/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/spark340/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.spark340.SparkSqlUtils
 import com.intel.raydp.shims.{ShimDescriptor, SparkShims}
 import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.StructType
 
 class Spark340Shims extends SparkShims {
@@ -49,4 +50,8 @@ class Spark340Shims extends SparkShims {
   override def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     SparkSqlUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  override def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]] = {
+    dataFrame.toArrowBatchRdd
+  }
 }
diff --git a/core/shims/spark340/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala b/core/shims/spark340/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.spark340
 import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.TaskContext
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, SQLContext, SparkSession}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types._
@@ -42,4 +43,8 @@ object SparkSqlUtils {
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     ArrowUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  def toArrowRDD(dataFrame: DataFrame, sparkSession: SparkSession): RDD[Array[Byte]] = {
+    SparkSqlUtils.toArrowRDD(dataFrame, dataFrame.sparkSession)
+  }
 }
diff --git a/core/shims/spark350/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/spark350/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.spark350.SparkSqlUtils
 import com.intel.raydp.shims.{ShimDescriptor, SparkShims}
 import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types.StructType
 
 class Spark350Shims extends SparkShims {
@@ -49,4 +50,8 @@ class Spark350Shims extends SparkShims {
   override def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     SparkSqlUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  override def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]] = {
+    SparkSqlUtils.toArrowRDD(dataFrame, dataFrame.sparkSession)
+  }
 }
diff --git a/core/shims/spark350/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala b/core/shims/spark350/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.spark350
 import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.TaskContext
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, SQLContext, SparkSession}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types._
@@ -42,4 +43,8 @@ object SparkSqlUtils {
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     ArrowUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId, errorOnDuplicatedFieldNames = false)
   }
+
+  def toArrowRDD(dataFrame: DataFrame, sparkSession: SparkSession): RDD[Array[Byte]] = {
+    dataFrame.toArrowBatchRdd
+  }
 }
diff --git a/core/shims/spark400/src/main/scala/com/intel/raydp/shims/SparkShims.scala b/core/shims/spark400/src/main/scala/com/intel/raydp/shims/SparkShims.scala
@@ -22,10 +22,15 @@ import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.executor.{RayDPExecutorBackendFactory, RayDPSpark400ExecutorBackendFactory}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.spark400.SparkSqlUtils
-import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.spark400.TaskContextUtils
 import com.intel.raydp.shims.{ShimDescriptor, SparkShims}
+import org.apache.spark.rdd.{MapPartitionsRDD, RDD}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.arrow.ArrowConverters
+
+import scala.reflect.ClassTag
 
 class Spark400Shims extends SparkShims {
   override def getShimDescriptor: ShimDescriptor = SparkShimProvider.DESCRIPTOR
@@ -48,4 +53,8 @@ class Spark400Shims extends SparkShims {
   override def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {
     SparkSqlUtils.toArrowSchema(schema = schema, timeZoneId = timeZoneId)
   }
+
+  override def toArrowBatchRDD(dataFrame: DataFrame): RDD[Array[Byte]] = {
+    SparkSqlUtils.toArrowRDD(dataFrame, dataFrame.sparkSession)
+  }
 }
diff --git a/core/shims/spark400/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala b/core/shims/spark400/src/main/scala/org/apache/spark/sql/SparkSqlUtils.scala
@@ -20,8 +20,9 @@ package org.apache.spark.sql.spark400
 import org.apache.arrow.vector.types.pojo.Schema
 import org.apache.spark.TaskContext
 import org.apache.spark.api.java.JavaRDD
-import org.apache.spark.sql.classic.ClassicConversions
-import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.classic.ClassicConversions.castToImpl
+import org.apache.spark.sql.{DataFrame, Row, SparkSession}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.util.ArrowUtils
@@ -43,7 +44,11 @@ object SparkSqlUtils {
         largeVarTypes = false,
         context = context)
     }
-    ClassicConversions.castToImpl(session).internalCreateDataFrame(internalRowRdd.setName("arrow"), schema)
+    session.internalCreateDataFrame(internalRowRdd.setName("arrow"), schema)
+  }
+
+  def toArrowRDD(dataFrame: DataFrame, sparkSession: SparkSession): RDD[Array[Byte]] = {
+    dataFrame.toArrowBatchRdd
   }
 
   def toArrowSchema(schema : StructType, timeZoneId : String) : Schema = {

Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ object ObjectStoreWriter {`
`225`	`225`	`}`
`226`	`226`	`val uuid = dfToId.getOrElseUpdate(df, UUID.randomUUID())`
`227`	`227`	`val queue = ObjectRefHolder.getQueue(uuid)`
`228`		`- val rdd = df.toArrowBatchRdd`
	`228`	`+ val rdd = SparkShimLoader.getSparkShims.toArrowBatchRDD(df)`
`229`	`229`	`rdd.persist(storageLevel)`
`230`	`230`	`rdd.count()`
`231`	`231`	`var executorIds = df.sqlContext.sparkContext.getExecutorIds.toArray`