This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch branch-0.4
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/branch-0.4 by this push:
new ef6eabb00 [CELEBORN-1300] Optimize CelebornInputStreamImpl's memory
usage
ef6eabb00 is described below
commit ef6eabb00e6fe020eee111bb817d920e58ba699c
Author: zky.zhoukeyong <[email protected]>
AuthorDate: Tue Mar 5 14:03:11 2024 +0800
[CELEBORN-1300] Optimize CelebornInputStreamImpl's memory usage
### What changes were proposed in this pull request?
To avoid too much memory usage when CelebornShuffleReader creates input
streams.
This PR does the following:
1. Constructor of `CelebornInputStream` does not fetch chunk
2. `compressedBuf` and `rawDataBuf` are created first time `fillBuffer` is
called
3. When `fillBuffer` returns false, which means the inputstream is
exhausted, `close` is called and resource released
4. `CelebornFetchFailureSuite` is only run for Spark 3.0 and newer
### Why are the changes needed?
ditto
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
GA and e2e test.
Closes #2348 from waitinfuture/1300.
Lead-authored-by: zky.zhoukeyong <[email protected]>
Co-authored-by: Keyong Zhou <[email protected]>
Signed-off-by: zky.zhoukeyong <[email protected]>
(cherry picked from commit 8b6bc35997a97a4252fc4c95dac4ca0b6d15aba8)
Signed-off-by: zky.zhoukeyong <[email protected]>
---
.../celeborn/client/read/CelebornInputStream.java | 101 +++---
.../tests/spark/CelebornFetchFailureSuite.scala | 346 +++++++++++----------
.../celeborn/tests/spark/SparkTestBase.scala | 5 +
3 files changed, 248 insertions(+), 204 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index 95a5ad5b2..a10239bf5 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -124,19 +124,20 @@ public abstract class CelebornInputStream extends
InputStream {
private final CelebornConf conf;
private final TransportClientFactory clientFactory;
private final String shuffleKey;
- private final PartitionLocation[] locations;
- private final int[] attempts;
+ private PartitionLocation[] locations;
+ private int[] attempts;
private final int attemptNumber;
private final int startMapIndex;
private final int endMapIndex;
- private final Map<Integer, Set<Integer>> batchesRead = new HashMap<>();
+ private Map<Integer, Set<Integer>> batchesRead = new HashMap<>();
private byte[] compressedBuf;
private byte[] rawDataBuf;
private Decompressor decompressor;
private ByteBuf currentChunk;
+ private boolean firstChunk = true;
private PartitionReader currentReader;
private final int fetchChunkMaxRetry;
private int fetchChunkRetryCnt = 0;
@@ -159,7 +160,7 @@ public abstract class CelebornInputStream extends
InputStream {
private boolean fetchExcludeWorkerOnFailureEnabled;
private boolean shuffleCompressionEnabled;
private long fetchExcludedWorkerExpireTimeout;
- private final ConcurrentHashMap<String, Long> fetchExcludedWorkers;
+ private ConcurrentHashMap<String, Long> fetchExcludedWorkers;
private boolean containLocalRead = false;
private ShuffleClient shuffleClient;
@@ -167,6 +168,7 @@ public abstract class CelebornInputStream extends
InputStream {
private int shuffleId;
private int partitionId;
private ExceptionMaker exceptionMaker;
+ private boolean closed = false;
CelebornInputStreamImpl(
CelebornConf conf,
@@ -203,16 +205,6 @@ public abstract class CelebornInputStream extends
InputStream {
this.fetchExcludedWorkerExpireTimeout =
conf.clientFetchExcludedWorkerExpireTimeout();
this.fetchExcludedWorkers = fetchExcludedWorkers;
- int bufferSize = conf.clientFetchBufferSize();
- if (shuffleCompressionEnabled) {
- int headerLen = Decompressor.getCompressionHeaderLength(conf);
- bufferSize += headerLen;
- compressedBuf = new byte[bufferSize];
-
- decompressor = Decompressor.getDecompressor(conf);
- }
- rawDataBuf = new byte[bufferSize];
-
if (conf.clientPushReplicateEnabled()) {
fetchChunkMaxRetry = conf.clientFetchMaxRetriesForEachReplica() * 2;
} else {
@@ -228,7 +220,7 @@ public abstract class CelebornInputStream extends
InputStream {
this.shuffleId = shuffleId;
this.shuffleClient = shuffleClient;
- moveToNextReader();
+ moveToNextReader(false);
}
private boolean skipLocation(int startMapIndex, int endMapIndex,
PartitionLocation location) {
@@ -270,7 +262,7 @@ public abstract class CelebornInputStream extends
InputStream {
return currentLocation;
}
- private void moveToNextReader() throws IOException {
+ private void moveToNextReader(boolean fetchChunk) throws IOException {
if (currentReader != null) {
currentReader.close();
currentReader = null;
@@ -291,7 +283,9 @@ public abstract class CelebornInputStream extends
InputStream {
currentReader = createReaderWithRetry(currentLocation);
fileIndex++;
}
- currentChunk = getNextChunk();
+ if (fetchChunk) {
+ currentChunk = getNextChunk();
+ }
}
private void excludeFailedLocation(PartitionLocation location, Exception
e) {
@@ -517,25 +511,40 @@ public abstract class CelebornInputStream extends
InputStream {
}
@Override
- public void close() {
- int locationsCount = locations.length;
- logger.debug(
- "total location count {} read {} skip {}",
- locationsCount,
- locationsCount - skipCount.sum(),
- skipCount.sum());
- if (currentChunk != null) {
- logger.debug("Release chunk {}", currentChunk);
- currentChunk.release();
- currentChunk = null;
- }
- if (currentReader != null) {
- logger.debug("Closing reader");
- currentReader.close();
- currentReader = null;
- }
- if (containLocalRead) {
- ShuffleClient.printReadStats(logger);
+ public synchronized void close() {
+ if (!closed) {
+ int locationsCount = locations.length;
+ logger.debug(
+ "AppShuffleId {}, shuffleId {}, partitionId {}, total location
count {}, read {}, skip {}",
+ appShuffleId,
+ shuffleId,
+ partitionId,
+ locationsCount,
+ locationsCount - skipCount.sum(),
+ skipCount.sum());
+ if (currentChunk != null) {
+ logger.debug("Release chunk {}", currentChunk);
+ currentChunk.release();
+ currentChunk = null;
+ }
+ if (currentReader != null) {
+ logger.debug("Closing reader");
+ currentReader.close();
+ currentReader = null;
+ }
+ if (containLocalRead) {
+ ShuffleClient.printReadStats(logger);
+ }
+
+ compressedBuf = null;
+ rawDataBuf = null;
+ batchesRead = null;
+ locations = null;
+ attempts = null;
+ decompressor = null;
+ fetchExcludedWorkers = null;
+
+ closed = true;
}
}
@@ -548,7 +557,7 @@ public abstract class CelebornInputStream extends
InputStream {
currentChunk = getNextChunk();
return true;
} else if (fileIndex < locations.length) {
- moveToNextReader();
+ moveToNextReader(true);
return currentReader != null;
}
if (currentReader != null) {
@@ -558,9 +567,27 @@ public abstract class CelebornInputStream extends
InputStream {
return false;
}
+ private void init() {
+ int bufferSize = conf.clientFetchBufferSize();
+
+ if (shuffleCompressionEnabled) {
+ int headerLen = Decompressor.getCompressionHeaderLength(conf);
+ bufferSize += headerLen;
+ compressedBuf = new byte[bufferSize];
+ decompressor = Decompressor.getDecompressor(conf);
+ }
+ rawDataBuf = new byte[bufferSize];
+ }
+
private boolean fillBuffer() throws IOException {
try {
+ if (firstChunk && currentReader != null) {
+ init();
+ currentChunk = getNextChunk();
+ firstChunk = false;
+ }
if (currentChunk == null) {
+ close();
return false;
}
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala
index 291bdca18..bdc834b5c 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala
@@ -93,193 +93,205 @@ class CelebornFetchFailureSuite extends AnyFunSuite
}
test("celeborn spark integration test - Fetch Failure") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
- .config(
- "spark.shuffle.manager",
- "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
- .getOrCreate()
-
- val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
- val hook = new ShuffleReaderGetHook(celebornConf)
- TestCelebornShuffleManager.registerReaderGetHook(hook)
-
- val value = Range(1, 10000).mkString(",")
- val tuples = sparkSession.sparkContext.parallelize(1 to 10000, 2)
- .map { i => (i, value) }.groupByKey(16).collect()
-
- // verify result
- assert(hook.executed.get() == true)
- assert(tuples.length == 10000)
- for (elem <- tuples) {
- assert(elem._2.mkString(",").equals(value))
- }
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
+ .config(
+ "spark.shuffle.manager",
+ "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
+ .getOrCreate()
+
+ val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
+ val hook = new ShuffleReaderGetHook(celebornConf)
+ TestCelebornShuffleManager.registerReaderGetHook(hook)
+
+ val value = Range(1, 10000).mkString(",")
+ val tuples = sparkSession.sparkContext.parallelize(1 to 10000, 2)
+ .map { i => (i, value) }.groupByKey(16).collect()
+
+ // verify result
+ assert(hook.executed.get() == true)
+ assert(tuples.length == 10000)
+ for (elem <- tuples) {
+ assert(elem._2.mkString(",").equals(value))
+ }
- val shuffleMgr = SparkContextHelper.env
- .shuffleManager
- .asInstanceOf[TestCelebornShuffleManager]
- val lifecycleManager = shuffleMgr.getLifecycleManager
+ val shuffleMgr = SparkContextHelper.env
+ .shuffleManager
+ .asInstanceOf[TestCelebornShuffleManager]
+ val lifecycleManager = shuffleMgr.getLifecycleManager
- shuffleMgr.unregisterShuffle(0)
- assert(lifecycleManager.getUnregisterShuffleTime().containsKey(0))
- assert(lifecycleManager.getUnregisterShuffleTime().containsKey(1))
+ shuffleMgr.unregisterShuffle(0)
+ assert(lifecycleManager.getUnregisterShuffleTime().containsKey(0))
+ assert(lifecycleManager.getUnregisterShuffleTime().containsKey(1))
- sparkSession.stop()
+ sparkSession.stop()
+ }
}
test("celeborn spark integration test - unregister shuffle with
throwsFetchFailure disabled") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "false")
- .getOrCreate()
-
- val value = Range(1, 10000).mkString(",")
- val tuples = sparkSession.sparkContext.parallelize(1 to 10000, 2)
- .map { i => (i, value) }.groupByKey(16).collect()
-
- // verify result
- assert(tuples.length == 10000)
- for (elem <- tuples) {
- assert(elem._2.mkString(",").equals(value))
- }
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure",
"false")
+ .getOrCreate()
+
+ val value = Range(1, 10000).mkString(",")
+ val tuples = sparkSession.sparkContext.parallelize(1 to 10000, 2)
+ .map { i => (i, value) }.groupByKey(16).collect()
+
+ // verify result
+ assert(tuples.length == 10000)
+ for (elem <- tuples) {
+ assert(elem._2.mkString(",").equals(value))
+ }
- val shuffleMgr = SparkContextHelper.env
- .shuffleManager
- .asInstanceOf[SparkShuffleManager]
- val lifecycleManager = shuffleMgr.getLifecycleManager
+ val shuffleMgr = SparkContextHelper.env
+ .shuffleManager
+ .asInstanceOf[SparkShuffleManager]
+ val lifecycleManager = shuffleMgr.getLifecycleManager
- shuffleMgr.unregisterShuffle(0)
- assert(lifecycleManager.getUnregisterShuffleTime().containsKey(0))
+ shuffleMgr.unregisterShuffle(0)
+ assert(lifecycleManager.getUnregisterShuffleTime().containsKey(0))
- sparkSession.stop()
+ sparkSession.stop()
+ }
}
test("celeborn spark integration test - Fetch Failure with multiple shuffle
data") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
- .config(
- "spark.shuffle.manager",
- "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
- .getOrCreate()
-
- val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
- val hook = new ShuffleReaderGetHook(celebornConf)
- TestCelebornShuffleManager.registerReaderGetHook(hook)
-
- import sparkSession.implicits._
-
- val df1 = Seq((1, "a"), (2, "b")).toDF("id", "data").groupBy("id").count()
- val df2 = Seq((2, "c"), (2, "d")).toDF("id", "data").groupBy("id").count()
- val tuples = df1.hint("merge").join(df2, "id").select("*").collect()
-
- // verify result
- assert(hook.executed.get() == true)
- val expect = "[2,1,2]"
- assert(tuples.head.toString().equals(expect))
- sparkSession.stop()
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
+ .config(
+ "spark.shuffle.manager",
+ "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
+ .getOrCreate()
+
+ val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
+ val hook = new ShuffleReaderGetHook(celebornConf)
+ TestCelebornShuffleManager.registerReaderGetHook(hook)
+
+ import sparkSession.implicits._
+
+ val df1 = Seq((1, "a"), (2, "b")).toDF("id",
"data").groupBy("id").count()
+ val df2 = Seq((2, "c"), (2, "d")).toDF("id",
"data").groupBy("id").count()
+ val tuples = df1.hint("merge").join(df2, "id").select("*").collect()
+
+ // verify result
+ assert(hook.executed.get() == true)
+ val expect = "[2,1,2]"
+ assert(tuples.head.toString().equals(expect))
+ sparkSession.stop()
+ }
}
test("celeborn spark integration test - Fetch Failure with RDD reuse") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
- .config(
- "spark.shuffle.manager",
- "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
- .getOrCreate()
-
- val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
- val hook = new ShuffleReaderGetHook(celebornConf)
- TestCelebornShuffleManager.registerReaderGetHook(hook)
-
- val sc = sparkSession.sparkContext
- val rdd1 = sc.parallelize(0 until 10000, 3).map(v => (v, v)).groupByKey()
- val rdd2 = sc.parallelize(0 until 10000, 2).map(v => (v, v)).groupByKey()
- val rdd3 = rdd1.map(v => (v._2, v._1))
-
- hook.executed.set(true)
-
- rdd1.count()
- rdd2.count()
-
- hook.executed.set(false)
- rdd3.count()
- hook.executed.set(false)
- rdd3.count()
- hook.executed.set(false)
- rdd3.count()
- hook.executed.set(false)
- rdd3.count()
-
- sparkSession.stop()
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
+ .config(
+ "spark.shuffle.manager",
+ "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
+ .getOrCreate()
+
+ val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
+ val hook = new ShuffleReaderGetHook(celebornConf)
+ TestCelebornShuffleManager.registerReaderGetHook(hook)
+
+ val sc = sparkSession.sparkContext
+ val rdd1 = sc.parallelize(0 until 10000, 3).map(v => (v, v)).groupByKey()
+ val rdd2 = sc.parallelize(0 until 10000, 2).map(v => (v, v)).groupByKey()
+ val rdd3 = rdd1.map(v => (v._2, v._1))
+
+ hook.executed.set(true)
+
+ rdd1.count()
+ rdd2.count()
+
+ hook.executed.set(false)
+ rdd3.count()
+ hook.executed.set(false)
+ rdd3.count()
+ hook.executed.set(false)
+ rdd3.count()
+ hook.executed.set(false)
+ rdd3.count()
+
+ sparkSession.stop()
+ }
}
test("celeborn spark integration test - Fetch Failure with read write
shuffles in one stage") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
- .config(
- "spark.shuffle.manager",
- "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
- .getOrCreate()
-
- val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
- val hook = new ShuffleReaderGetHook(celebornConf)
- TestCelebornShuffleManager.registerReaderGetHook(hook)
-
- val sc = sparkSession.sparkContext
- val rdd1 = sc.parallelize(0 until 10000, 3).map(v => (v, v)).groupByKey()
- val rdd2 = rdd1.map(v => (v._2, v._1)).groupByKey()
-
- hook.executed.set(true)
- rdd1.count()
-
- hook.executed.set(false)
- rdd2.count()
-
- sparkSession.stop()
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
+ .config(
+ "spark.shuffle.manager",
+ "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
+ .getOrCreate()
+
+ val celebornConf =
SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
+ val hook = new ShuffleReaderGetHook(celebornConf)
+ TestCelebornShuffleManager.registerReaderGetHook(hook)
+
+ val sc = sparkSession.sparkContext
+ val rdd1 = sc.parallelize(0 until 10000, 3).map(v => (v, v)).groupByKey()
+ val rdd2 = rdd1.map(v => (v._2, v._1)).groupByKey()
+
+ hook.executed.set(true)
+ rdd1.count()
+
+ hook.executed.set(false)
+ rdd2.count()
+
+ sparkSession.stop()
+ }
}
test("celeborn spark integration test - empty shuffle data") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
- val sparkSession = SparkSession.builder()
- .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
- .config("spark.sql.shuffle.partitions", 2)
- .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
- .config("spark.celeborn.shuffle.enabled", "true")
- .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
- .getOrCreate()
-
- sparkSession.sql("create table if not exists t_1 (a bigint) using parquet")
- sparkSession.sql("create table if not exists t_2 (a bigint) using parquet")
- sparkSession.sql("create table if not exists t_3 (a bigint) using parquet")
- val df1 = sparkSession.table("t_1")
- val df2 = sparkSession.table("t_2")
- val df3 = sparkSession.table("t_3")
- df1.union(df2).union(df3).count()
-
- sparkSession.stop()
+ if (Spark3OrNewer) {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+ val sparkSession = SparkSession.builder()
+ .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+ .config("spark.sql.shuffle.partitions", 2)
+ .config("spark.celeborn.shuffle.forceFallback.partition.enabled",
false)
+ .config("spark.celeborn.shuffle.enabled", "true")
+ .config("spark.celeborn.client.spark.fetch.throwsFetchFailure", "true")
+ .getOrCreate()
+
+ sparkSession.sql("create table if not exists t_1 (a bigint) using
parquet")
+ sparkSession.sql("create table if not exists t_2 (a bigint) using
parquet")
+ sparkSession.sql("create table if not exists t_3 (a bigint) using
parquet")
+ val df1 = sparkSession.table("t_1")
+ val df2 = sparkSession.table("t_2")
+ val df3 = sparkSession.table("t_3")
+ df1.union(df2).union(df3).count()
+
+ sparkSession.stop()
+ }
}
}
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
index e2cb3c98a..05af928e6 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
@@ -19,6 +19,7 @@ package org.apache.celeborn.tests.spark
import scala.util.Random
+import org.apache.spark.SPARK_VERSION
import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.internal.SQLConf
@@ -32,6 +33,10 @@ import org.apache.celeborn.service.deploy.MiniClusterFeature
trait SparkTestBase extends AnyFunSuite
with Logging with MiniClusterFeature with BeforeAndAfterAll with
BeforeAndAfterEach {
+
+ val Spark3OrNewer = SPARK_VERSION >= "3.0"
+ println(s"Spark version is $SPARK_VERSION, Spark3OrNewer: $Spark3OrNewer")
+
private val sampleSeq = (1 to 78)
.map(Random.alphanumeric)
.toList
