This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new d007e293 [AURON #1372]add spark.auron.enable.shuffleExchange to
control ShuffleExchangeExec into Native or not (#1373)
d007e293 is described below
commit d007e2931610711aa1343a79f18b77d4df4edc96
Author: guixiaowen <[email protected]>
AuthorDate: Sat Oct 11 14:25:35 2025 +0800
[AURON #1372]add spark.auron.enable.shuffleExchange to control
ShuffleExchangeExec into Native or not (#1373)
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set #1373
* [AURON #1372]Do not convert into ShuffleExchangeExec into Native when
shuffle manger not set #1373
* [AURON #1378]add spark.auron.enable.broadcastexchange to control
BroadcastExchangeExec into Native or not #1378
* [AURON #1372]add spark.auron.enable.shuffleexchange to control
ShuffleExchangeExec into Native or not #1373
* [AURON #1372]add spark.auron.enable.shuffleExchange to control
ShuffleExchangeExec into Native or not #1373
---------
Co-authored-by: guihuawen <[email protected]>
---
.../AuronCheckConvertShuffleExchangeSuite.scala | 96 ++++++++++++++++++++++
.../apache/spark/sql/auron/AuronConverters.scala | 14 +++-
2 files changed, 109 insertions(+), 1 deletion(-)
diff --git
a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/auron/AuronCheckConvertShuffleExchangeSuite.scala
b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/auron/AuronCheckConvertShuffleExchangeSuite.scala
new file mode 100644
index 00000000..1d303064
--- /dev/null
+++
b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/auron/AuronCheckConvertShuffleExchangeSuite.scala
@@ -0,0 +1,96 @@
+/*
+ * 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.auron
+
+import org.apache.spark.sql.{QueryTest, Row, SparkSession}
+import org.apache.spark.sql.execution.auron.plan.NativeShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.test.SharedSparkSession
+
+class AuronCheckConvertShuffleExchangeSuite
+ extends QueryTest
+ with SharedSparkSession
+ with AuronSQLTestHelper
+ with org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper {
+
+ test(
+ "test set auron shuffle manager convert to native shuffle exchange where
set spark.auron.enable is true") {
+ withTable("test_shuffle") {
+ val spark = SparkSession
+ .builder()
+ .master("local[2]")
+ .appName("checkConvertToNativeShuffleManger")
+ .config("spark.sql.shuffle.partitions", "4")
+ .config("spark.sql.autoBroadcastJoinThreshold", -1)
+ .config("spark.sql.extensions",
"org.apache.spark.sql.auron.AuronSparkSessionExtension")
+ .config(
+ "spark.shuffle.manager",
+ "org.apache.spark.sql.execution.auron.shuffle.AuronShuffleManager")
+ .config("spark.memory.offHeap.enabled", "false")
+ .config("spark.auron.enable", "true")
+ .getOrCreate()
+
+ spark.sql("drop table if exists test_shuffle")
+ spark.sql(
+ "create table if not exists test_shuffle using parquet PARTITIONED BY
(part) as select 1 as c1, 2 as c2, 'test test' as part")
+ val executePlan =
+ spark.sql("select c1, count(1) from test_shuffle group by c1")
+
+ val shuffleExchangeExec =
+ executePlan.queryExecution.executedPlan
+ .collectFirst { case shuffleExchangeExec: ShuffleExchangeExec =>
+ shuffleExchangeExec
+ }
+ val afterConvertPlan =
AuronConverters.convertSparkPlan(shuffleExchangeExec.get)
+ assert(afterConvertPlan.isInstanceOf[NativeShuffleExchangeExec])
+ checkAnswer(executePlan, Seq(Row(1, 1)))
+ }
+ }
+
+ test(
+ "test set non auron shuffle manager do not convert to native shuffle
exchange where set spark.auron.enable is true") {
+ withTable("test_shuffle") {
+ val spark = SparkSession
+ .builder()
+ .master("local[2]")
+ .appName("checkConvertToNativeShuffleManger")
+ .config("spark.sql.shuffle.partitions", "4")
+ .config("spark.sql.autoBroadcastJoinThreshold", -1)
+ .config("spark.shuffle.manager",
"org.apache.spark.shuffle.sort.SortShuffleManager")
+ .config("spark.sql.extensions",
"org.apache.spark.sql.auron.AuronSparkSessionExtension")
+ .config("spark.memory.offHeap.enabled", "false")
+ .config("spark.auron.enable", "true")
+ .getOrCreate()
+ spark.sql("drop table if exists test_shuffle")
+ spark.sql(
+ "create table if not exists test_shuffle using parquet PARTITIONED BY
(part) as select 1 as c1, 2 as c2, 'test test' as part")
+ val executePlan =
+ spark.sql("select c1, count(1) from test_shuffle group by c1")
+
+ val shuffleExchangeExec =
+ executePlan.queryExecution.executedPlan
+ .collectFirst { case shuffleExchangeExec: ShuffleExchangeExec =>
+ shuffleExchangeExec
+ }
+ val afterConvertPlan =
AuronConverters.convertSparkPlan(shuffleExchangeExec.get)
+ assert(afterConvertPlan.isInstanceOf[ShuffleExchangeExec])
+ checkAnswer(executePlan, Seq(Row(1, 1)))
+
+ }
+ }
+
+}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala
index df2608b4..556ac7f3 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala
@@ -26,7 +26,7 @@ import org.apache.commons.lang3.reflect.MethodUtils
import org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat
import org.apache.spark.Partition
import org.apache.spark.broadcast.Broadcast
-import org.apache.spark.internal.Logging
+import org.apache.spark.internal.{config, Logging}
import
org.apache.spark.sql.auron.AuronConvertStrategy.{childOrderingRequiredTag,
convertibleTag, convertStrategyTag, convertToNonNativeTag, isNeverConvert,
joinSmallerSideTag, neverConvertReasonTag}
import org.apache.spark.sql.auron.NativeConverters.{roundRobinTypeSupported,
scalarTypeSupported, StubExpr}
import org.apache.spark.sql.auron.util.AuronLogUtils.logDebugPlanConversion
@@ -135,12 +135,21 @@ object AuronConverters extends Logging {
getBooleanConf("spark.auron.enable.scan.orc", defaultValue = true)
def enableBroadcastExchange: Boolean =
getBooleanConf("spark.auron.enable.broadcastExchange", defaultValue = true)
+ def enableShuffleExechange: Boolean =
+ getBooleanConf("spark.auron.enable.shuffleExchange", defaultValue = true)
private val extConvertProviders =
ServiceLoader.load(classOf[AuronConvertProvider]).asScala
def extConvertSupported(exec: SparkPlan): Boolean = {
extConvertProviders.exists(_.isSupported(exec))
}
+ def enableExchange(): Boolean = {
+ val shuffleMangerName =
SQLConf.get.getConfString(config.SHUFFLE_MANAGER.key)
+ enableShuffleExechange && !shuffleMangerName.isEmpty &&
(shuffleMangerName.contains(
+ "AuronShuffleManager") || shuffleMangerName.contains(
+ "AuronUniffleShuffleManager") ||
shuffleMangerName.contains("AuronCelebornShuffleManager"))
+ }
+
// format: off
// scalafix:off
// necessary imports for cross spark versions build
@@ -177,6 +186,9 @@ object AuronConverters extends Logging {
case e: ShuffleExchangeExec => tryConvert(e, convertShuffleExchangeExec)
case e: BroadcastExchangeExec if enableBroadcastExchange =>
tryConvert(e, convertBroadcastExchangeExec)
+ case e: ShuffleExchangeExec if enableExchange => tryConvert(e,
convertShuffleExchangeExec)
+ case e: BroadcastExchangeExec =>
+ tryConvert(e, convertBroadcastExchangeExec)
case e: FileSourceScanExec if enableScan => // scan
tryConvert(e, convertFileSourceScanExec)
case e: ProjectExec if enableProject => // project
