add option

Author: zuston

client-spark/common/pom.xml

@@ -95,6 +95,14 @@
             <artifactId>fory-core</artifactId>
             <version>0.12.0</version>
         </dependency>
+        
+        <!-- Scala dependencies -->
+        <dependency>
+            <groupId>org.scala-lang</groupId>
+            <artifactId>scala-library</artifactId>
+            <version>${scala.version}</version>
+            <scope>provided</scope>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -119,8 +127,17 @@
                 <version>${scala.maven.plugin.version}</version>
                 <executions>
                     <execution>
+                        <id>scala-compile-first</id>
+                        <phase>process-resources</phase>
                         <goals>
+                            <goal>add-source</goal>
                             <goal>compile</goal>
+                        </goals>
+                    </execution>
+                    <execution>
+                        <id>scala-test-compile-first</id>
+                        <phase>process-test-resources</phase>
+                        <goals>
                             <goal>testCompile</goal>
                         </goals>
                     </execution>
@@ -129,6 +146,19 @@
                     <scalaVersion>${scala.version}</scalaVersion>
                 </configuration>
             </plugin>
+            
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <phase>compile</phase>
+                        <goals>
+                            <goal>compile</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
 
             <!-- ScalaTest Maven plugin for running tests -->
             <plugin>

client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java

@@ -36,6 +36,7 @@
 import org.apache.uniffle.common.config.ConfigUtils;
 import org.apache.uniffle.common.config.RssClientConf;
 import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.shuffle.ShuffleSerializer;
 
 public class RssSparkConfig {
 
@@ -110,6 +111,12 @@ public class RssSparkConfig {
           .defaultValue(true)
           .withDescription("indicates row based shuffle, set false when use in columnar shuffle");
 
+  public static final ConfigOption<ShuffleSerializer> RSS_SHUFFLE_SERIALIZER =
+      ConfigOptions.key("rss.client.shuffle.serializer")
+          .enumType(ShuffleSerializer.class)
+          .noDefaultValue()
+          .withDescription("Shuffle serializer type");
+
   public static final ConfigOption<Boolean> RSS_MEMORY_SPILL_ENABLED =
       ConfigOptions.key("rss.client.memory.spill.enabled")
           .booleanType()

client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java

@@ -40,9 +40,9 @@
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.serializer.ForySerializerInstance;
 import org.apache.spark.serializer.SerializationStream;
 import org.apache.spark.serializer.Serializer;
-import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.shuffle.RssSparkConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -58,6 +58,7 @@
 import org.apache.uniffle.common.util.BlockIdLayout;
 import org.apache.uniffle.common.util.ChecksumUtils;
 
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_SHUFFLE_SERIALIZER;
 import static org.apache.spark.shuffle.RssSparkConfig.RSS_WRITE_OVERLAPPING_COMPRESSION_ENABLED;
 
 public class WriteBufferManager extends MemoryConsumer {
@@ -81,7 +82,6 @@ public class WriteBufferManager extends MemoryConsumer {
   private int shuffleId;
   private String taskId;
   private long taskAttemptId;
-  private SerializerInstance instance;
   private ShuffleWriteMetrics shuffleWriteMetrics;
   // cache partition -> records
   private Map<Integer, WriterBuffer> buffers;
@@ -192,8 +192,11 @@ public WriteBufferManager(
     // in columnar shuffle, the serializer here is never used
     this.isRowBased = rssConf.getBoolean(RssSparkConfig.RSS_ROW_BASED);
     if (isRowBased) {
-      this.instance = serializer.newInstance();
-      this.serializeStream = instance.serializeStream(arrayOutputStream);
+      if (rssConf.contains(RSS_SHUFFLE_SERIALIZER)) {
+        this.serializeStream = new ForySerializerInstance().serializeStream(arrayOutputStream);
+      } else {
+        this.serializeStream = serializer.newInstance().serializeStream(arrayOutputStream);
+      }
     }
     boolean compress =
         rssConf.getBoolean(

client-spark/common/src/main/java/org/apache/uniffle/shuffle/ShuffleSerializer.java

@@ -0,0 +1,22 @@
+/*
+ * 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.uniffle.shuffle;
+
+public enum ShuffleSerializer {
+  FORY
+}

client-spark/common/src/main/scala/org/apache/spark/serializer/ForySerializer.scala

@@ -31,15 +31,6 @@ class ForySerializerInstance extends org.apache.spark.serializer.SerializerInsta
     f
   })
 
-  private def createFuryInstance(): Fory = {
-    Fory.builder()
-      .withLanguage(Language.JAVA)
-      .withRefTracking(true)
-      .withCompatibleMode(CompatibleMode.COMPATIBLE)
-      .requireClassRegistration(false)
-      .build()
-  }
-
   override def serialize[T: ClassTag](t: T): ByteBuffer = {
     val bytes = fury.get().serialize(t.asInstanceOf[AnyRef])
     ByteBuffer.wrap(bytes)

client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java

@@ -42,6 +42,7 @@
 import org.apache.spark.ShuffleDependency;
 import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleReadMetrics;
+import org.apache.spark.serializer.ForySerializer;
 import org.apache.spark.serializer.Serializer;
 import org.apache.spark.shuffle.FunctionUtils;
 import org.apache.spark.shuffle.RssShuffleHandle;
@@ -68,6 +69,7 @@
 
 import static org.apache.spark.shuffle.RssSparkConfig.RSS_READ_REORDER_MULTI_SERVERS_ENABLED;
 import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_SHUFFLE_SERIALIZER;
 
 public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
   private static final Logger LOG = LoggerFactory.getLogger(RssShuffleReader.class);
@@ -125,7 +127,10 @@ public RssShuffleReader(
     this.numMaps = rssShuffleHandle.getNumMaps();
     this.shuffleDependency = rssShuffleHandle.getDependency();
     this.shuffleId = shuffleDependency.shuffleId();
-    this.serializer = rssShuffleHandle.getDependency().serializer();
+    this.serializer =
+        rssConf.contains(RSS_SHUFFLE_SERIALIZER)
+            ? new ForySerializer()
+            : rssShuffleHandle.getDependency().serializer();
     this.taskId = "" + context.taskAttemptId() + "_" + context.attemptNumber();
     this.basePath = basePath;
     this.partitionNum = partitionNum;