You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/27 20:17:20 UTC
(datafusion-comet) branch main updated: feat: add Celeborn shuffle manager and partition pusher (5/n) (#5501)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new f7629904f8 feat: add Celeborn shuffle manager and partition pusher (5/n) (#5501)
f7629904f8 is described below
commit f7629904f818e30fcb87a80e1c3773eeddc5e85b
Author: Ping Zhang <pi...@gmail.com>
AuthorDate: Thu Aug 27 13:17:15 2026 -0700
feat: add Celeborn shuffle manager and partition pusher (5/n) (#5501)
* feat: add Celeborn shuffle manager and partition pusher
* fix: preserve Celeborn integrity accounting and encrypted pushes
---
.github/workflows/pr_build_linux.yml | 2 +
.github/workflows/pr_build_macos.yml | 2 +
.../shuffle/CelebornShufflePartitionPusher.java | 236 +++++++++
.../shuffle/CometCelebornShuffleManager.scala | 145 ++++++
.../CelebornShufflePartitionPusherSuite.scala | 536 +++++++++++++++++++++
.../shuffle/CometCelebornShuffleManagerSuite.scala | 282 +++++++++++
6 files changed, 1203 insertions(+)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index b7c88254c9..6b7130b401 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -327,6 +327,8 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+ org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 8210f91b7f..65d87ec24f 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -143,6 +143,8 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+ org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
new file mode 100644
index 0000000000..132961ac44
--- /dev/null
+++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
@@ -0,0 +1,236 @@
+/*
+ * 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.comet.shuffle;
+
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/** Adapts complete Comet shuffle frames to an existing, task-owned Celeborn shuffle client. */
+public final class CelebornShufflePartitionPusher implements ShufflePartitionPusher {
+
+ private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES;
+ private static final int MINIMUM_COMET_FRAME_BYTES = 2 * Long.BYTES;
+
+ private final Object shuffleClient;
+ private final Method pushOrMergeData;
+ private final Method computeBatchCRC;
+ private final int shuffleId;
+ private final int mapId;
+ private final int encodedAttemptId;
+ private final int numMappers;
+ private final int numPartitions;
+
+ /**
+ * Captures all shuffle and map-attempt identity before native worker threads invoke this pusher.
+ *
+ * <p>Celeborn remains an optional, application-provided dependency: accepting its client as an
+ * {@link Object} and resolving its public raw-push method avoids a compile-time dependency.
+ *
+ * @param shuffleClient the application's existing Celeborn shuffle client
+ * @param shuffleId the Celeborn shuffle identifier for this task
+ * @param mapId the logical map partition
+ * @param encodedAttemptId the Celeborn-encoded stage and task attempt
+ * @param numMappers the total number of map partitions
+ * @param numPartitions the total number of reduce partitions
+ */
+ public CelebornShufflePartitionPusher(
+ Object shuffleClient,
+ int shuffleId,
+ int mapId,
+ int encodedAttemptId,
+ int numMappers,
+ int numPartitions) {
+ if (shuffleClient == null) {
+ throw new IllegalArgumentException("Celeborn shuffle client must not be null");
+ }
+ if (shuffleId < 0) {
+ throw new IllegalArgumentException("Celeborn shuffle ID must not be negative");
+ }
+ if (mapId < 0) {
+ throw new IllegalArgumentException("Celeborn map ID must not be negative");
+ }
+ if (encodedAttemptId < 0) {
+ throw new IllegalArgumentException("Celeborn encoded attempt ID must not be negative");
+ }
+ if (numMappers <= 0) {
+ throw new IllegalArgumentException("Celeborn mapper count must be positive");
+ }
+ if (mapId >= numMappers) {
+ throw new IllegalArgumentException("Celeborn map ID is outside the mapper count");
+ }
+ if (numPartitions <= 0) {
+ throw new IllegalArgumentException("Celeborn partition count must be positive");
+ }
+
+ final Method pushMethod;
+ try {
+ pushMethod =
+ shuffleClient
+ .getClass()
+ .getMethod(
+ "pushOrMergeData",
+ int.class,
+ int.class,
+ int.class,
+ int.class,
+ byte[].class,
+ int.class,
+ int.class,
+ int.class,
+ int.class,
+ boolean.class,
+ boolean.class);
+ } catch (NoSuchMethodException | SecurityException e) {
+ throw new IllegalArgumentException(
+ "Celeborn shuffle client does not provide the required public raw-push API", e);
+ }
+
+ if (pushMethod.getReturnType() != int.class || Modifier.isStatic(pushMethod.getModifiers())) {
+ throw new IllegalArgumentException(
+ "Celeborn raw-push API must be an instance method returning an int");
+ }
+
+ Method integrityMethod = null;
+ try {
+ integrityMethod =
+ shuffleClient
+ .getClass()
+ .getMethod(
+ "computeBatchCRC",
+ int.class,
+ int.class,
+ int.class,
+ int.class,
+ byte[].class,
+ int.class,
+ int.class);
+ } catch (NoSuchMethodException missing) {
+ // Older Celeborn clients account for integrity inside pushOrMergeData instead.
+ try {
+ for (Method method : shuffleClient.getClass().getMethods()) {
+ if (method.getName().equals("computeBatchCRC")) {
+ throw new IllegalArgumentException(
+ "Celeborn integrity-accounting API has an incompatible signature", missing);
+ }
+ }
+ } catch (SecurityException failure) {
+ throw new IllegalArgumentException(
+ "Cannot inspect the optional Celeborn integrity-accounting API", failure);
+ }
+ } catch (SecurityException failure) {
+ throw new IllegalArgumentException(
+ "Cannot resolve the optional Celeborn integrity-accounting API", failure);
+ }
+
+ if (integrityMethod != null
+ && (integrityMethod.getReturnType() != void.class
+ || Modifier.isStatic(integrityMethod.getModifiers()))) {
+ throw new IllegalArgumentException(
+ "Celeborn integrity-accounting API must be an instance method returning void");
+ }
+
+ this.shuffleClient = shuffleClient;
+ this.pushOrMergeData = pushMethod;
+ this.computeBatchCRC = integrityMethod;
+ this.shuffleId = shuffleId;
+ this.mapId = mapId;
+ this.encodedAttemptId = encodedAttemptId;
+ this.numMappers = numMappers;
+ this.numPartitions = numPartitions;
+ }
+
+ /** Sends exactly one complete, already-compressed Comet shuffle frame to Celeborn. */
+ @Override
+ public void pushPartitionData(int partitionId, byte[] data, int length) throws IOException {
+ if (partitionId < 0 || partitionId >= numPartitions) {
+ throw new IOException("Celeborn output partition is outside this task's partition count");
+ }
+ if (data == null) {
+ throw new IOException("Celeborn shuffle frame must not be null");
+ }
+ if (length > Integer.MAX_VALUE - CELEBORN_BATCH_HEADER_BYTES) {
+ throw new IOException("Celeborn shuffle frame and transport header exceed the byte limit");
+ }
+ if (length < MINIMUM_COMET_FRAME_BYTES || length > data.length) {
+ throw new IOException("Celeborn shuffle frame length must describe one complete frame");
+ }
+
+ final long declaredBodyLength = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong();
+ if (declaredBodyLength != (long) length - Long.BYTES) {
+ throw new IOException(
+ "Celeborn shuffle frame declares "
+ + declaredBodyLength
+ + " body bytes, but contains "
+ + (length - Long.BYTES));
+ }
+
+ final int accepted;
+ try {
+ if (computeBatchCRC != null) {
+ computeBatchCRC.invoke(
+ shuffleClient, shuffleId, mapId, encodedAttemptId, partitionId, data, 0, length);
+ }
+
+ accepted =
+ (int)
+ pushOrMergeData.invoke(
+ shuffleClient,
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ partitionId,
+ data,
+ 0,
+ length,
+ numMappers,
+ numPartitions,
+ true,
+ true);
+ } catch (IllegalAccessException e) {
+ throw new IOException("Cannot invoke the public Celeborn raw-push API", e);
+ } catch (InvocationTargetException e) {
+ Throwable cause = e.getCause();
+ if (cause instanceof IOException) {
+ throw (IOException) cause;
+ }
+ if (cause instanceof RuntimeException) {
+ throw (RuntimeException) cause;
+ }
+ if (cause instanceof Error) {
+ throw (Error) cause;
+ }
+ throw new IOException("Celeborn raw shuffle push failed", cause);
+ }
+
+ int minimumAccepted = length + CELEBORN_BATCH_HEADER_BYTES;
+ if (accepted < minimumAccepted) {
+ throw new IOException(
+ "Celeborn raw shuffle push accepted "
+ + accepted
+ + " bytes; expected at least "
+ + minimumAccepted
+ + " including its transport header");
+ }
+ }
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
new file mode 100644
index 0000000000..a3759df11e
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.execution.shuffle
+
+import java.lang.reflect.InvocationTargetException
+
+import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+
+import org.apache.comet.util.ClassLoaders
+
+/**
+ * Preserves the application's existing Celeborn shuffle manager for ordinary Spark shuffles.
+ *
+ * Celeborn is an optional, application-provided dependency, so its shuffle manager is loaded
+ * reflectively. Comet shuffle dependencies remain unsupported until subsequent changes add their
+ * remote map-side lifecycle and reduce-side reader.
+ */
+class CometCelebornShuffleManager private[shuffle] (
+ conf: SparkConf,
+ isDriver: Boolean,
+ backendFactory: (SparkConf, Boolean) => ShuffleManager)
+ extends ShuffleManager {
+
+ /** Constructor used by Spark on both the driver and executors. */
+ def this(conf: SparkConf, isDriver: Boolean) =
+ this(conf, isDriver, CometCelebornShuffleManager.createBackend)
+
+ private val backend = Option(backendFactory(conf, isDriver)).getOrElse {
+ throw new IllegalStateException("Celeborn Spark shuffle manager factory returned null")
+ }
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ dependency match {
+ case _: CometShuffleDependency[_, _, _] => rejectCometShuffle()
+ case _ => backend.registerShuffle(shuffleId, dependency)
+ }
+ }
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ rejectCometHandle(handle)
+ backend.getWriter(handle, mapId, context, metrics)
+ }
+
+ // Spark's final all-mapper reader overload delegates to this mapper-range overload.
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ rejectCometHandle(handle)
+ backend.getReader(
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ }
+
+ override def shuffleBlockResolver: ShuffleBlockResolver = backend.shuffleBlockResolver
+
+ override def unregisterShuffle(shuffleId: Int): Boolean = backend.unregisterShuffle(shuffleId)
+
+ override def stop(): Unit = backend.stop()
+
+ private def rejectCometHandle(handle: ShuffleHandle): Unit = handle match {
+ case _: CometNativeShuffleHandle[_, _] => rejectCometShuffle()
+ case _: CometBypassMergeSortShuffleHandle[_, _] => rejectCometShuffle()
+ case _: CometSerializedShuffleHandle[_, _] => rejectCometShuffle()
+ case base: BaseShuffleHandle[_, _, _]
+ if base.dependency.isInstanceOf[CometShuffleDependency[_, _, _]] =>
+ rejectCometShuffle()
+ case _ =>
+ }
+
+ private def rejectCometShuffle(): Nothing = {
+ throw new UnsupportedOperationException(
+ "Comet shuffle over Celeborn is not supported until native shuffle integration is complete")
+ }
+}
+
+private[shuffle] object CometCelebornShuffleManager {
+
+ private val CELEBORN_MANAGER_CLASS = "org.apache.spark.shuffle.celeborn.SparkShuffleManager"
+
+ private[shuffle] def createBackend(conf: SparkConf, isDriver: Boolean): ShuffleManager = {
+ try {
+ val managerClass = ClassLoaders.loadClass(CELEBORN_MANAGER_CLASS)
+ if (!classOf[ShuffleManager].isAssignableFrom(managerClass)) {
+ throw new IllegalStateException(
+ "Celeborn Spark shuffle manager does not implement ShuffleManager: " +
+ CELEBORN_MANAGER_CLASS)
+ }
+
+ val constructor = managerClass.getConstructor(classOf[SparkConf], java.lang.Boolean.TYPE)
+ constructor.newInstance(conf, Boolean.box(isDriver)).asInstanceOf[ShuffleManager]
+ } catch {
+ case failure: ClassNotFoundException =>
+ throw new IllegalStateException(
+ s"Celeborn Spark shuffle manager is not available: $CELEBORN_MANAGER_CLASS. " +
+ "Ensure the Celeborn Spark client is present on the application classpath",
+ failure)
+ case failure: InvocationTargetException =>
+ throw new IllegalStateException(
+ s"Could not initialize Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS",
+ Option(failure.getCause).getOrElse(failure))
+ case failure: ReflectiveOperationException =>
+ throw new IllegalStateException(
+ s"Could not construct Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS",
+ failure)
+ case failure: LinkageError =>
+ throw new IllegalStateException(
+ s"Could not load Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS",
+ failure)
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala
new file mode 100644
index 0000000000..3a661ecbc9
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala
@@ -0,0 +1,536 @@
+/*
+ * 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.comet.shuffle
+
+import java.io.IOException
+import java.nio.{ByteBuffer, ByteOrder}
+import java.util.Arrays
+import java.util.concurrent.atomic.AtomicReference
+import java.util.zip.CRC32
+
+import scala.collection.mutable
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class CelebornShufflePartitionPusherSuite extends AnyFunSuite {
+
+ private val encodedAttemptId = (4 << 16) | 7
+
+ private def pusher(client: AnyRef): CelebornShufflePartitionPusher =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9)
+
+ private def frame(bodyLength: Int = java.lang.Long.BYTES): Array[Byte] = {
+ val buffer = ByteBuffer
+ .allocate(java.lang.Long.BYTES + bodyLength)
+ .order(ByteOrder.LITTLE_ENDIAN)
+ .putLong(bodyLength.toLong)
+ (0 until bodyLength).foreach(index => buffer.put((index + 1).toByte))
+ buffer.array()
+ }
+
+ private def combinedCrc(frames: Array[Byte]*): Long = {
+ frames.foldLeft(0L) { (combined, bytes) =>
+ val checksum = new CRC32
+ checksum.update(bytes, 0, bytes.length)
+ (0 until java.lang.Integer.BYTES).foldLeft(0L) { (result, index) =>
+ val shift = index * java.lang.Byte.SIZE
+ val next = ((combined >>> shift) & 0xffL) + ((checksum.getValue >>> shift) & 0xffL)
+ result | ((next & 0xffL) << shift)
+ }
+ }
+ }
+
+ test("raw push forwards captured task metadata and preserves the complete Comet frame") {
+ val client = new RecordingCelebornPushClient
+ val bytes = frame()
+
+ pusher(client).pushPartitionData(6, bytes, bytes.length)
+
+ val push = client.lastPush
+ assert(push.shuffleId == 19)
+ assert(push.mapId == 3)
+ assert(push.attemptId == encodedAttemptId)
+ assert(push.partitionId == 6)
+ assert(push.bytes eq bytes)
+ assert(push.offset == 0)
+ assert(push.length == bytes.length)
+ assert(push.numMappers == 12)
+ assert(push.numPartitions == 9)
+ assert(push.doPush)
+ assert(push.skipCompress)
+ }
+
+ test("raw push accepts a complete frame within a larger backing array") {
+ val client = new RecordingCelebornPushClient
+ val completeFrame = frame()
+ val backingArray = Arrays.copyOf(completeFrame, completeFrame.length + 4)
+
+ pusher(client).pushPartitionData(0, backingArray, completeFrame.length)
+
+ assert(client.lastPush.bytes eq backingArray)
+ assert(client.lastPush.length == completeFrame.length)
+ }
+
+ test("raw push rejects fewer bytes than the frame and Celeborn transport header require") {
+ val bytes = frame()
+
+ Seq(0, bytes.length, bytes.length + 15, -1).foreach { accepted =>
+ val client = new RecordingCelebornPushClient
+ client.acceptedBytes = Some(accepted)
+
+ val failure = intercept[IOException] {
+ pusher(client).pushPartitionData(0, bytes, bytes.length)
+ }
+
+ assert(failure.getMessage.contains(accepted.toString))
+ assert(failure.getMessage.contains((bytes.length + 16).toString))
+ }
+ }
+
+ test("legacy Celeborn clients without integrity accounting remain compatible") {
+ val client = new RecordingCelebornPushClient
+ val bytes = frame()
+
+ assert(!client.getClass.getMethods.exists(_.getName == "computeBatchCRC"))
+ pusher(client).pushPartitionData(2, bytes, bytes.length)
+
+ assert(client.pushCount == 1)
+ assert(client.lastPush.partitionId == 2)
+ assert(client.lastPush.skipCompress)
+ }
+
+ test("integrity accounting records every complete frame exactly once before each raw push") {
+ val client = new IntegrityCheckingCelebornPushClient
+ val adapter = pusher(client)
+ val first = frame(8)
+ val second = frame(11)
+ val third = frame(12)
+
+ adapter.pushPartitionData(2, first, first.length)
+ adapter.pushPartitionData(5, second, second.length)
+ adapter.pushPartitionData(2, third, third.length)
+
+ assert(client.pushCount == 3)
+ assert(client.accountedFrames.size == 3)
+ assert(client.partitionCrc(2) == combinedCrc(first, third))
+ assert(client.partitionBytes(2) == first.length + third.length)
+ assert(client.partitionCrc(5) == combinedCrc(second))
+ assert(client.partitionBytes(5) == second.length)
+ assert(
+ client.invocationOrder.toSeq ==
+ Seq("crc:2", "push:2", "crc:5", "push:5", "crc:2", "push:2"))
+ assert(client.accountedFrames.forall(_.shuffleId == 19))
+ assert(client.accountedFrames.forall(_.mapId == 3))
+ assert(client.accountedFrames.forall(_.attemptId == encodedAttemptId))
+ assert(client.accountedFrames.forall(_.offset == 0))
+ assert(
+ client.accountedFrames
+ .map(_.length)
+ .toSeq == Seq(first.length, second.length, third.length))
+ assert(client.accountedFrames.head.bytes eq first)
+ assert(client.recordedPushes.forall(_.doPush))
+ assert(client.recordedPushes.forall(_.skipCompress))
+ }
+
+ test("integrity accounting failures prevent raw pushes and preserve the original exception") {
+ val bytes = frame()
+
+ Seq[Throwable](
+ new IOException("Celeborn integrity accounting rejected the frame"),
+ new IllegalStateException("Celeborn integrity accounting state was unavailable"))
+ .foreach { expected =>
+ val client = new IntegrityCheckingCelebornPushClient
+ client.integrityFailure = expected
+
+ val actual = intercept[Throwable] {
+ pusher(client).pushPartitionData(0, bytes, bytes.length)
+ }
+
+ assert(actual eq expected)
+ assert(client.pushCount == 0)
+ assert(client.accountedFrames.isEmpty)
+ assert(client.invocationOrder.toSeq == Seq("crc:0"))
+ }
+ }
+
+ test("raw push accepts frames expanded by Spark IO encryption without pushing twice") {
+ val client = new RecordingCelebornPushClient
+ val cryptoHandler = new RecordingCelebornCryptoHandler
+ val bytes = frame()
+ client.cryptoHandler = Some(cryptoHandler)
+
+ pusher(client).pushPartitionData(6, bytes, bytes.length)
+
+ assert(client.pushCount == 1)
+ assert(cryptoHandler.encryptionCount == 1)
+ assert(cryptoHandler.plaintext.sameElements(bytes))
+ assert(cryptoHandler.encryptedLength == bytes.length + 20)
+ assert(client.lastPush.bytes eq bytes)
+ assert(client.lastPush.skipCompress)
+ }
+
+ test("integrity accounting covers plaintext before Spark IO encryption expands the frame") {
+ val client = new IntegrityCheckingCelebornPushClient
+ val cryptoHandler = new RecordingCelebornCryptoHandler
+ val bytes = frame(13)
+ client.cryptoHandler = Some(cryptoHandler)
+
+ pusher(client).pushPartitionData(4, bytes, bytes.length)
+
+ assert(client.partitionCrc(4) == combinedCrc(bytes))
+ assert(client.partitionBytes(4) == bytes.length)
+ assert(client.invocationOrder.toSeq == Seq("crc:4", "push:4"))
+ assert(cryptoHandler.plaintext.sameElements(bytes))
+ assert(cryptoHandler.encryptedLength == bytes.length + 20)
+ assert(client.pushCount == 1)
+ }
+
+ test("raw push preserves the exact IOException thrown by the Celeborn client") {
+ val client = new RecordingCelebornPushClient
+ val expected = new IOException("Celeborn worker rejected the shuffle frame")
+ val bytes = frame()
+ client.failure = expected
+
+ val actual = intercept[IOException] {
+ pusher(client).pushPartitionData(0, bytes, bytes.length)
+ }
+
+ assert(actual eq expected)
+ }
+
+ test("raw push preserves unchecked exceptions and errors thrown by the Celeborn client") {
+ val bytes = frame()
+
+ val runtimeClient = new RecordingCelebornPushClient
+ val runtimeFailure = new IllegalStateException("Celeborn client was closed")
+ runtimeClient.failure = runtimeFailure
+ val actualRuntimeFailure = intercept[IllegalStateException] {
+ pusher(runtimeClient).pushPartitionData(0, bytes, bytes.length)
+ }
+ assert(actualRuntimeFailure eq runtimeFailure)
+
+ val errorClient = new RecordingCelebornPushClient
+ val expectedError = new AssertionError("Celeborn client invariant failed")
+ errorClient.failure = expectedError
+ val actualError = intercept[AssertionError] {
+ pusher(errorClient).pushPartitionData(0, bytes, bytes.length)
+ }
+ assert(actualError eq expectedError)
+ }
+
+ test("raw push wraps unexpected checked client failures without losing their cause") {
+ val client = new RecordingCelebornPushClient
+ val expected = new InterruptedException("Celeborn client was interrupted")
+ val bytes = frame()
+ client.failure = expected
+
+ val actual = intercept[IOException] {
+ pusher(client).pushPartitionData(0, bytes, bytes.length)
+ }
+
+ assert(actual.getCause eq expected)
+ }
+
+ test("adapter rejects a missing client or an incompatible public Celeborn raw-push API") {
+ intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(null, 0, 0, 0, 1, 1)
+ }
+
+ val missing = intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(new Object, 0, 0, 0, 1, 1)
+ }
+ assert(missing.getMessage.contains("public raw-push"))
+
+ val wrongReturnType = intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(new WrongReturnTypeCelebornPushClient, 0, 0, 0, 1, 1)
+ }
+ assert(wrongReturnType.getMessage.contains("returning an int"))
+ }
+
+ test("adapter rejects invalid shuffle, map-attempt, mapper, and partition metadata") {
+ val client = new RecordingCelebornPushClient
+ val invalidMetadata = Seq(
+ (-1, 0, 0, 1, 1),
+ (0, -1, 0, 1, 1),
+ (0, 0, -1, 1, 1),
+ (0, 0, 0, 0, 1),
+ (0, 1, 0, 1, 1),
+ (0, 0, 0, 1, 0))
+
+ invalidMetadata.foreach { case (shuffleId, mapId, attemptId, numMappers, numPartitions) =>
+ intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(
+ client,
+ shuffleId,
+ mapId,
+ attemptId,
+ numMappers,
+ numPartitions)
+ }
+ }
+ }
+
+ test("adapter rejects invalid output partitions and frame lengths before invoking Celeborn") {
+ val client = new RecordingCelebornPushClient
+ val adapter = pusher(client)
+ val bytes = frame()
+
+ Seq(-1, 9).foreach { partitionId =>
+ intercept[IOException] {
+ adapter.pushPartitionData(partitionId, bytes, bytes.length)
+ }
+ }
+
+ intercept[IOException] {
+ adapter.pushPartitionData(0, null, bytes.length)
+ }
+
+ Seq(-1, 0, java.lang.Long.BYTES, bytes.length - 1, bytes.length + 1).foreach { length =>
+ intercept[IOException] {
+ adapter.pushPartitionData(0, bytes, length)
+ }
+ }
+
+ val overflow = intercept[IOException] {
+ adapter.pushPartitionData(0, bytes, Int.MaxValue)
+ }
+ assert(overflow.getMessage.contains("transport header"))
+ assert(client.pushCount == 0)
+ }
+
+ test("adapter rejects truncated, concatenated, and incorrectly encoded shuffle frames") {
+ val client = new RecordingCelebornPushClient
+ val adapter = pusher(client)
+
+ Seq(-1L, 0L, 7L, 9L, Long.MaxValue).foreach { declaredLength =>
+ val bytes = frame()
+ ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).putLong(declaredLength)
+
+ val failure = intercept[IOException] {
+ adapter.pushPartitionData(0, bytes, bytes.length)
+ }
+ assert(failure.getMessage.contains("body bytes"))
+ }
+
+ val wrongByteOrder = frame()
+ ByteBuffer.wrap(wrongByteOrder).order(ByteOrder.BIG_ENDIAN).putLong(8L)
+ intercept[IOException] {
+ adapter.pushPartitionData(0, wrongByteOrder, wrongByteOrder.length)
+ }
+
+ assert(client.pushCount == 0)
+ }
+
+ test("captured map-task metadata remains available when native invokes a worker thread") {
+ val client = new RecordingCelebornPushClient
+ val adapter = pusher(client)
+ val bytes = frame()
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(new Runnable {
+ override def run(): Unit = {
+ try adapter.pushPartitionData(4, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ }
+ })
+
+ worker.start()
+ worker.join(5000)
+
+ assert(!worker.isAlive)
+ assert(failure.get() == null)
+ assert(client.lastPush.shuffleId == 19)
+ assert(client.lastPush.mapId == 3)
+ assert(client.lastPush.attemptId == encodedAttemptId)
+ assert(client.lastPush.partitionId == 4)
+ }
+}
+
+/** Public so the adapter can resolve and invoke the optional client's API using reflection. */
+class RecordingCelebornPushClient {
+
+ @volatile var acceptedBytes: Option[Int] = None
+ @volatile var cryptoHandler: Option[RecordingCelebornCryptoHandler] = None
+ @volatile var failure: Throwable = _
+ @volatile var lastPush: RecordedCelebornPush = _
+ @volatile var pushCount: Int = 0
+
+ @throws[IOException]
+ def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int = {
+ pushCount += 1
+ lastPush = RecordedCelebornPush(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+
+ if (failure != null) {
+ throw failure
+ }
+ val transportPayloadLength =
+ cryptoHandler.fold(length)(handler => handler.encrypt(bytes, offset, length).length)
+ acceptedBytes.getOrElse(transportPayloadLength + 16)
+ }
+}
+
+/** Models the Spark crypto wire format's minimum 4-byte length plus 16-byte IV overhead. */
+final class RecordingCelebornCryptoHandler {
+
+ @volatile var encryptionCount: Int = 0
+ @volatile var plaintext: Array[Byte] = _
+ @volatile var encryptedLength: Int = 0
+
+ def encrypt(bytes: Array[Byte], offset: Int, length: Int): Array[Byte] = {
+ encryptionCount += 1
+ plaintext = Arrays.copyOfRange(bytes, offset, offset + length)
+ encryptedLength = length + java.lang.Integer.BYTES + 16
+ new Array[Byte](encryptedLength)
+ }
+}
+
+/** Mirrors Celeborn 0.7 integrity accounting without depending on its optional client classes. */
+final class IntegrityCheckingCelebornPushClient extends RecordingCelebornPushClient {
+
+ @volatile var integrityFailure: Throwable = _
+ val accountedFrames: mutable.ArrayBuffer[RecordedCelebornAccounting] =
+ mutable.ArrayBuffer.empty
+ val recordedPushes: mutable.ArrayBuffer[RecordedCelebornPush] = mutable.ArrayBuffer.empty
+ val invocationOrder: mutable.ArrayBuffer[String] = mutable.ArrayBuffer.empty
+ private val checksums = mutable.HashMap.empty[Int, Long]
+ private val byteTotals = mutable.HashMap.empty[Int, Long]
+
+ def partitionCrc(partitionId: Int): Long = checksums(partitionId)
+
+ def partitionBytes(partitionId: Int): Long = byteTotals(partitionId)
+
+ @throws[IOException]
+ def computeBatchCRC(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int): Unit = {
+ invocationOrder += s"crc:$partitionId"
+ if (integrityFailure != null) {
+ throw integrityFailure
+ }
+
+ accountedFrames +=
+ RecordedCelebornAccounting(shuffleId, mapId, attemptId, partitionId, bytes, offset, length)
+ val batchChecksum = new CRC32
+ batchChecksum.update(bytes, offset, length)
+ val previous = checksums.getOrElse(partitionId, 0L)
+ val combined = (0 until java.lang.Integer.BYTES).foldLeft(0L) { (result, index) =>
+ val shift = index * java.lang.Byte.SIZE
+ val next = ((previous >>> shift) & 0xffL) + ((batchChecksum.getValue >>> shift) & 0xffL)
+ result | ((next & 0xffL) << shift)
+ }
+ checksums.update(partitionId, combined)
+ byteTotals.update(partitionId, byteTotals.getOrElse(partitionId, 0L) + length)
+ }
+
+ @throws[IOException]
+ override def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int = {
+ invocationOrder += s"push:$partitionId"
+ val accepted = super.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+ recordedPushes += lastPush
+ accepted
+ }
+}
+
+final case class RecordedCelebornAccounting(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int)
+
+final case class RecordedCelebornPush(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean)
+
+/** Mimics an incompatible optional client whose raw-push method does not return an int. */
+final class WrongReturnTypeCelebornPushClient {
+
+ def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Long = length.toLong
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
new file mode 100644
index 0000000000..486aa1e481
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
@@ -0,0 +1,282 @@
+/*
+ * 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.comet.execution.shuffle
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+
+class CometCelebornShuffleManagerSuite extends AnyFunSuite {
+
+ private class RecordingShuffleManager extends ShuffleManager {
+
+ val returnedHandle: ShuffleHandle = new BaseShuffleHandle[Any, Any, Any](31, null)
+
+ var registration: Option[(Int, ShuffleDependency[_, _, _])] = None
+ var writerCall: Option[(ShuffleHandle, Long)] = None
+ var rangedReaderCall: Option[(ShuffleHandle, Int, Int, Int, Int)] = None
+ var unregisteredShuffleId: Option[Int] = None
+ var unregisterResult = true
+ var resolverReads = 0
+ var stopped = false
+ var registrationFailure: RuntimeException = _
+ var writerFailure: RuntimeException = _
+ var readerFailure: RuntimeException = _
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ registration = Some((shuffleId, dependency))
+ if (registrationFailure != null) {
+ throw registrationFailure
+ }
+ returnedHandle
+ }
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ writerCall = Some((handle, mapId))
+ if (writerFailure != null) {
+ throw writerFailure
+ }
+ null
+ }
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ rangedReaderCall = Some((handle, startMapIndex, endMapIndex, startPartition, endPartition))
+ if (readerFailure != null) {
+ throw readerFailure
+ }
+ null
+ }
+
+ override def shuffleBlockResolver: ShuffleBlockResolver = {
+ resolverReads += 1
+ null
+ }
+
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ unregisteredShuffleId = Some(shuffleId)
+ unregisterResult
+ }
+
+ override def stop(): Unit = stopped = true
+ }
+
+ private def manager(
+ backend: ShuffleManager,
+ conf: SparkConf = new SparkConf(false),
+ isDriver: Boolean = true): CometCelebornShuffleManager = {
+ new CometCelebornShuffleManager(conf, isDriver, (_, _) => backend)
+ }
+
+ test("manager preserves the application Spark configuration and driver identity") {
+ val conf = new SparkConf(false).set("spark.app.name", "existing-celeborn-application")
+ val backend = new RecordingShuffleManager
+ var observedConf: SparkConf = null
+ var observedIsDriver = true
+
+ new CometCelebornShuffleManager(
+ conf,
+ false,
+ (actualConf, actualIsDriver) => {
+ observedConf = actualConf
+ observedIsDriver = actualIsDriver
+ backend
+ })
+
+ assert(observedConf eq conf)
+ assert(!observedIsDriver)
+ assert(conf.get("spark.app.name") == "existing-celeborn-application")
+ }
+
+ test("ordinary shuffle registration preserves the delegated manager's handle") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val dependency = null.asInstanceOf[ShuffleDependency[Any, Any, Any]]
+
+ val handle = composite.registerShuffle(31, dependency)
+
+ assert(handle eq backend.returnedHandle)
+ assert(backend.registration.contains((31, dependency)))
+ }
+
+ test("ordinary map writers are owned by the delegated Celeborn manager") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val handle = backend.returnedHandle
+
+ assert(composite.getWriter[Any, Any](handle, 93L, null, null) == null)
+ assert(backend.writerCall.contains((handle, 93L)))
+ }
+
+ test("mapper-range reads preserve the delegated Celeborn reader's exact range") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val handle = backend.returnedHandle
+
+ assert(composite.getReader[Any, Any](handle, 2, 8, 3, 7, null, null) == null)
+ assert(backend.rangedReaderCall.contains((handle, 2, 8, 3, 7)))
+ }
+
+ test("all-mapper reads delegate through Spark's inherited complete mapper range") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val handle = backend.returnedHandle
+
+ assert(composite.getReader[Any, Any](handle, 4, 9, null, null) == null)
+ assert(backend.rangedReaderCall.contains((handle, 0, Int.MaxValue, 4, 9)))
+ }
+
+ test("resolver, shuffle removal, and shutdown retain the delegated lifecycle") {
+ val backend = new RecordingShuffleManager
+ backend.unregisterResult = false
+ val composite = manager(backend)
+
+ assert(composite.shuffleBlockResolver == null)
+ assert(backend.resolverReads == 1)
+ assert(!composite.unregisterShuffle(17))
+ assert(backend.unregisteredShuffleId.contains(17))
+
+ composite.stop()
+
+ assert(backend.stopped)
+ }
+
+ test("registration failures retain their original exception without local fallback") {
+ val backend = new RecordingShuffleManager
+ val expected = new IllegalStateException("remote shuffle registration failed")
+ backend.registrationFailure = expected
+
+ val actual = intercept[IllegalStateException] {
+ manager(backend).registerShuffle[Any, Any, Any](31, null)
+ }
+
+ assert(actual eq expected)
+ assert(backend.writerCall.isEmpty)
+ assert(backend.rangedReaderCall.isEmpty)
+ }
+
+ test("delegated writer and reader failures retain their original exception") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val handle = backend.returnedHandle
+ val writerFailure = new IllegalStateException("remote shuffle writer failed")
+ val readerFailure = new IllegalStateException("remote shuffle reader failed")
+ backend.writerFailure = writerFailure
+ backend.readerFailure = readerFailure
+
+ val actualWriterFailure = intercept[IllegalStateException] {
+ composite.getWriter[Any, Any](handle, 0L, null, null)
+ }
+ val actualAllMapperFailure = intercept[IllegalStateException] {
+ composite.getReader[Any, Any](handle, 0, 1, null, null)
+ }
+ val actualRangedFailure = intercept[IllegalStateException] {
+ composite.getReader[Any, Any](handle, 0, 1, 0, 1, null, null)
+ }
+
+ assert(actualWriterFailure eq writerFailure)
+ assert(actualAllMapperFailure eq readerFailure)
+ assert(actualRangedFailure eq readerFailure)
+ }
+
+ test("Comet native and JVM handles never reach Celeborn's ordinary shuffle paths") {
+ val backend = new RecordingShuffleManager
+ val composite = manager(backend)
+ val unsupportedHandles = Seq[ShuffleHandle](
+ new CometNativeShuffleHandle[Any, Any](41, null),
+ new CometBypassMergeSortShuffleHandle[Any, Any](42, null),
+ new CometSerializedShuffleHandle[Any, Any](43, null))
+
+ unsupportedHandles.foreach { handle =>
+ val writerFailure = intercept[UnsupportedOperationException] {
+ composite.getWriter[Any, Any](handle, 0L, null, null)
+ }
+ val allMapperFailure = intercept[UnsupportedOperationException] {
+ composite.getReader[Any, Any](handle, 0, 1, null, null)
+ }
+ val rangedFailure = intercept[UnsupportedOperationException] {
+ composite.getReader[Any, Any](handle, 0, 1, 0, 1, null, null)
+ }
+
+ assert(writerFailure.getMessage.contains("Comet shuffle over Celeborn is not supported"))
+ assert(allMapperFailure.getMessage.contains("Comet shuffle over Celeborn is not supported"))
+ assert(rangedFailure.getMessage.contains("Comet shuffle over Celeborn is not supported"))
+ }
+
+ assert(backend.writerCall.isEmpty)
+ assert(backend.rangedReaderCall.isEmpty)
+ }
+
+ test("a null delegated backend fails without selecting a local shuffle manager") {
+ val failure = intercept[IllegalStateException] {
+ new CometCelebornShuffleManager(new SparkConf(false), true, (_, _) => null)
+ }
+
+ assert(failure.getMessage.contains("factory returned null"))
+ }
+
+ test("a delegated backend construction failure retains its original exception") {
+ val expected = new IllegalStateException("existing Celeborn manager failed to initialize")
+
+ val actual = intercept[IllegalStateException] {
+ new CometCelebornShuffleManager(new SparkConf(false), true, (_, _) => throw expected)
+ }
+
+ assert(actual eq expected)
+ }
+
+ test("the public constructor rejects an application without the optional Celeborn client") {
+ val currentThread = Thread.currentThread()
+ val originalLoader = currentThread.getContextClassLoader
+ val missingCelebornLoader = new ClassLoader(originalLoader) {
+ override protected def loadClass(name: String, resolve: Boolean): Class[_] = {
+ if (name == "org.apache.spark.shuffle.celeborn.SparkShuffleManager") {
+ throw new ClassNotFoundException(name)
+ }
+ super.loadClass(name, resolve)
+ }
+ }
+
+ currentThread.setContextClassLoader(missingCelebornLoader)
+ try {
+ val failure = intercept[IllegalStateException] {
+ new CometCelebornShuffleManager(new SparkConf(false), true)
+ }
+
+ assert(failure.getMessage.contains("Celeborn Spark shuffle manager is not available"))
+ assert(failure.getCause.isInstanceOf[ClassNotFoundException])
+ } finally {
+ currentThread.setContextClassLoader(originalLoader)
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org