From 714b4f103b6a3370772dee1665810e5682d94e96 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <liangchi@apple.com>
Date: Sat, 12 Aug 2023 08:43:46 -0700
Subject: [PATCH] rdar://113781872: Enable BosonShuffleManager in Spark 3.4
 (#1845)

---
 core/pom.xml                                     |  5 +++++
 .../main/scala/org/apache/spark/SparkEnv.scala   | 16 +++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/core/pom.xml b/core/pom.xml
index 799e0e4fe980f..5fe013dd85b6f 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -492,6 +492,11 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>com.apple.pie.boson</groupId>
+      <artifactId>boson-spark-spark${spark.version.short}_${scala.binary.version}</artifactId>
+    </dependency>
+
     <!--
       This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
       them will yield errors.
diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 28ab9dc074232..e15deebb90425 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -26,6 +26,7 @@ import scala.collection.concurrent
 import scala.collection.mutable
 import scala.util.Properties
 
+import com.apple.boson.BosonConf
 import com.google.common.cache.CacheBuilder
 import org.apache.hadoop.conf.Configuration
 
@@ -316,7 +317,20 @@ object SparkEnv extends Logging {
     val shortShuffleMgrNames = Map(
       "sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName,
       "tungsten-sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName)
-    val shuffleMgrName = conf.get(config.SHUFFLE_MANAGER)
+    val shuffleMgrName = {
+      val bosonEnv = System.getenv("BOSON")
+      val bosonEnabled =
+        conf.getBoolean(BosonConf.BOSON_ENABLED.key, bosonEnv == null || bosonEnv.toBoolean)
+      val bosonShuffleEnabled = conf.getBoolean(BosonConf.BOSON_EXEC_SHUFFLE_ENABLED.key, false)
+
+      if (!conf.contains(config.SHUFFLE_MANAGER) && bosonEnabled && bosonShuffleEnabled) {
+        // If no custom shuffle manager is specified, and Boson is enabled,
+        // use Boson shuffle manager.
+        "org.apache.spark.sql.boson.execution.shuffle.BosonShuffleManager"
+      } else {
+        conf.get(config.SHUFFLE_MANAGER)
+      }
+    }
     val shuffleMgrClass =
       shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase(Locale.ROOT), shuffleMgrName)
     val shuffleManager = Utils.instantiateSerializerOrShuffleManager[ShuffleManager](