diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 31cc3c34da9..11a5619a4d2 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -176,6 +176,49 @@ jobs:
- name: Compile (skip tests)
run: ./mvnw -B install -DskipTests -Dmaven.test.skip=true -Pspark-4.1
+ celeborn-reflection-compatibility:
+ needs: lint
+ name: Celeborn ${{ matrix.celeborn_version }} reflection compatibility
+ runs-on: ubuntu-24.04
+ container:
+ image: amd64/rust
+ env:
+ JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED
+ strategy:
+ fail-fast: false
+ matrix:
+ celeborn_version: ["0.6.0", "0.7.0"]
+ steps:
+ - uses: actions/checkout@v7
+
+ - name: Setup Rust & Java toolchain
+ uses: ./.github/actions/setup-builder
+ with:
+ rust-version: ${{ env.RUST_VERSION }}
+ jdk-version: 17
+
+ - name: Cache Maven dependencies
+ uses: actions/cache@v6
+ with:
+ path: |
+ ~/.m2/repository
+ /root/.m2/repository
+ key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-celeborn-${{ matrix.celeborn_version }}
+ restore-keys: |
+ ${{ runner.os }}-java-maven-
+
+ - name: Verify reflected Celeborn internals
+ env:
+ SPARK_LOCAL_HOSTNAME: localhost
+ SPARK_LOCAL_IP: 127.0.0.1
+ run: |
+ SPARK_HOME="$GITHUB_WORKSPACE" ./mvnw -B clean test \
+ -Pspark-3.5,scala-2.12,celeborn-reflection-compatibility \
+ -Dceleborn.version="${{ matrix.celeborn_version }}" \
+ -Dtest=none \
+ -Dsuites=org.apache.comet.shuffle.CelebornReflectionCompatibilitySuite \
+ -DfailIfNoTests=false
+
# Build native library once and share with all test jobs
build-native:
needs: lint
diff --git a/dev/ci/check-suites.py b/dev/ci/check-suites.py
index b7369d17076..52a221b2cd3 100644
--- a/dev/ci/check-suites.py
+++ b/dev/ci/check-suites.py
@@ -36,6 +36,7 @@ def file_to_class_name(path: Path) -> str | None:
"org.apache.comet.parquet.ParquetReadFromS3Suite", # manual test suite
"org.apache.comet.IcebergReadFromS3Suite", # manual test suite
"org.apache.comet.cloud.s3.CometS3CredentialBridgeSuite", # manual test suite
+ "org.apache.comet.shuffle.CelebornReflectionCompatibilitySuite", # dedicated version matrix
"org.apache.spark.sql.comet.CometPlanStabilitySuite", # abstract
"org.apache.spark.sql.comet.ParquetDatetimeRebaseSuite", # abstract
"org.apache.comet.exec.CometColumnarShuffleSuite" # abstract
diff --git a/pom.xml b/pom.xml
index 2b34cb4a25e..f4b2be220ec 100644
--- a/pom.xml
+++ b/pom.xml
@@ -52,6 +52,8 @@ under the License.
property without an error.
-->
4.1.0
+
+ 0.7.0
${java.version}
${java.version}
3.11.0
@@ -1046,6 +1048,7 @@ under the License.
src/main/scala/**/*.scala
src/test/scala/**/*.scala
+ src/test/celeborn-reflection-compatibility/**/*.scala
src/main/spark-*/**/*.scala
src/test/spark-*/**/*.scala
diff --git a/spark/pom.xml b/spark/pom.xml
index a257415dd3f..8dc632d5b92 100644
--- a/spark/pom.xml
+++ b/spark/pom.xml
@@ -222,6 +222,41 @@ under the License.
+
+ celeborn-reflection-compatibility
+
+
+
+ org.apache.celeborn
+ celeborn-client-spark-3-shaded_${scala.binary.version}
+ ${celeborn.version}
+ test
+
+
+
+
+
+ org.codehaus.mojo
+ build-helper-maven-plugin
+
+
+ add-celeborn-reflection-compatibility-test-source
+ generate-test-sources
+
+ add-test-source
+
+
+
+ src/test/celeborn-reflection-compatibility
+
+
+
+
+
+
+
+
+
spark-3.4
diff --git a/spark/src/test/celeborn-reflection-compatibility/org/apache/comet/shuffle/CelebornReflectionCompatibilitySuite.scala b/spark/src/test/celeborn-reflection-compatibility/org/apache/comet/shuffle/CelebornReflectionCompatibilitySuite.scala
new file mode 100644
index 00000000000..1830c2d5433
--- /dev/null
+++ b/spark/src/test/celeborn-reflection-compatibility/org/apache/comet/shuffle/CelebornReflectionCompatibilitySuite.scala
@@ -0,0 +1,221 @@
+/*
+ * 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.lang.reflect.{Field, Method, Modifier}
+import java.util.{List => JList, Map => JMap, Optional}
+import java.util.concurrent.{ConcurrentHashMap, ExecutorService}
+import java.util.concurrent.atomic.{AtomicReference, LongAdder}
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class CelebornReflectionCompatibilitySuite extends AnyFunSuite {
+
+ private val classLoader = getClass.getClassLoader
+
+ private lazy val shuffleClient = load("org.apache.celeborn.client.ShuffleClientImpl")
+ private lazy val pushState = load("org.apache.celeborn.common.write.PushState")
+ private lazy val inFlightRequestTracker =
+ load("org.apache.celeborn.common.write.InFlightRequestTracker")
+ private lazy val transportClientFactory =
+ load("org.apache.celeborn.common.network.client.TransportClientFactory")
+ private lazy val transportClient =
+ load("org.apache.celeborn.common.network.client.TransportClient")
+ private lazy val transportResponseHandler =
+ load("org.apache.celeborn.common.network.client.TransportResponseHandler")
+ private lazy val pushRequestInfo = load("org.apache.celeborn.common.write.PushRequestInfo")
+
+ private lazy val celebornVersion = Option(shuffleClient.getPackage)
+ .flatMap(pkg => Option(pkg.getImplementationVersion))
+ .getOrElse(fail("Celeborn client JAR does not declare its implementation version"))
+
+ private lazy val celebornReleaseLine = celebornVersion.split("\\.").take(2).mkString(".")
+
+ test("load a released Celeborn 0.6 or 0.7 client") {
+ assert(shuffleClient.getPackage.getImplementationTitle == "celeborn-client-spark-3-shaded")
+ assert(
+ Set("0.6", "0.7").contains(celebornReleaseLine),
+ s"unsupported Celeborn compatibility-test version $celebornVersion")
+ }
+
+ test("resolve the Celeborn shuffle client members used by the partition pusher") {
+ instanceMethod(
+ shuffleClient,
+ "pushOrMergeData",
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ classOf[Array[Byte]],
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Boolean.TYPE,
+ java.lang.Boolean.TYPE)
+ instanceMethod(
+ shuffleClient,
+ "cleanup",
+ java.lang.Void.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE)
+ instanceMethod(shuffleClient, "getPushState", pushState, classOf[String])
+
+ instanceField(shuffleClient, "pushStates", classOf[JMap[_, _]])
+ instanceField(shuffleClient, "pushDataRetryPool", classOf[ExecutorService])
+
+ celebornReleaseLine match {
+ case "0.6" =>
+ instanceMethod(
+ shuffleClient,
+ "mapperEnd",
+ java.lang.Void.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE)
+ assert(!shuffleClient.getMethods.exists(_.getName == "computeBatchCRC"))
+ intercept[NoSuchFieldException](declaredField(shuffleClient, "cryptoHandler"))
+
+ case "0.7" =>
+ instanceMethod(
+ shuffleClient,
+ "mapperEnd",
+ java.lang.Void.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE)
+ instanceMethod(
+ shuffleClient,
+ "computeBatchCRC",
+ java.lang.Void.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ classOf[Array[Byte]],
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE)
+ instanceField(shuffleClient, "cryptoHandler", classOf[Optional[_]])
+
+ case other => fail(s"unsupported Celeborn release line $other")
+ }
+ }
+
+ test("resolve the Celeborn push-state members used for admission") {
+ exactInstanceField(pushState, "exception", classOf[AtomicReference[_]])
+ instanceField(pushState, "inFlightRequestTracker", inFlightRequestTracker)
+ exactInstanceField(inFlightRequestTracker, "totalInflightReqs", classOf[LongAdder])
+ }
+
+ test("resolve the Celeborn transport members used to track callback ownership") {
+ instanceMethod(shuffleClient, "getDataClientFactory", transportClientFactory)
+
+ val transportClientBootstrap =
+ load("org.apache.celeborn.common.network.client.TransportClientBootstrap")
+ assert(transportClientBootstrap.isInterface)
+ instanceMethod(transportClientBootstrap, "doBootstrap", java.lang.Void.TYPE, transportClient)
+
+ val bootstraps = instanceField(transportClientFactory, "clientBootstraps", classOf[JList[_]])
+ assert(bootstraps.getGenericType.getTypeName.contains(transportClientBootstrap.getName))
+ instanceField(transportClientFactory, "connectionPool", classOf[JMap[_, _]])
+
+ val clientPool =
+ load("org.apache.celeborn.common.network.client.TransportClientFactory$ClientPool")
+ val clients = instanceField(clientPool, "clients", classOf[Array[_]])
+ assert(clients.getType.getComponentType == transportClient)
+ val locks = instanceField(clientPool, "locks", classOf[Array[_]])
+ assert(locks.getType.getComponentType == classOf[Object])
+
+ val channel = declaredField(transportClient, "channel")
+ assert(!Modifier.isStatic(channel.getModifiers))
+ assert(channel.getType.isInterface)
+ instanceMethod(transportClient, "getHandler", transportResponseHandler)
+ instanceField(transportResponseHandler, "outstandingPushes", classOf[ConcurrentHashMap[_, _]])
+
+ val callback = declaredField(pushRequestInfo, "callback")
+ assert(!Modifier.isStatic(callback.getModifiers))
+ assert(!Modifier.isFinal(callback.getModifiers))
+ assert(callback.getType.isInterface)
+
+ val writeFuture = channel.getType.getMethod("writeAndFlush", classOf[Object]).getReturnType
+ assert(writeFuture.isInterface)
+ instanceMethod(writeFuture, "isDone", java.lang.Boolean.TYPE)
+ assert(writeFuture.getMethods.exists { method =>
+ method.getName == "addListener" &&
+ method.getParameterCount == 1 &&
+ method.getParameterTypes.head.isInterface
+ })
+ }
+
+ private def load(name: String): Class[_] = Class.forName(name, false, classLoader)
+
+ private def instanceField(owner: Class[_], name: String, expectedType: Class[_]): Field = {
+ val result = declaredField(owner, name)
+ assert(
+ !Modifier.isStatic(result.getModifiers),
+ s"${owner.getName}.$name must be an instance field")
+ assert(
+ expectedType.isAssignableFrom(result.getType),
+ s"${owner.getName}.$name has type ${result.getType.getName}, expected ${expectedType.getName}")
+ result
+ }
+
+ private def exactInstanceField(owner: Class[_], name: String, expectedType: Class[_]): Field = {
+ val result = instanceField(owner, name, expectedType)
+ assert(
+ result.getType == expectedType,
+ s"${owner.getName}.$name has type ${result.getType.getName}, expected ${expectedType.getName}")
+ result
+ }
+
+ private def declaredField(owner: Class[_], name: String): Field = {
+ var current = owner
+ while (current != null) {
+ try {
+ val result = current.getDeclaredField(name)
+ result.setAccessible(true)
+ return result
+ } catch {
+ case _: NoSuchFieldException => current = current.getSuperclass
+ }
+ }
+ throw new NoSuchFieldException(s"${owner.getName}.$name")
+ }
+
+ private def instanceMethod(
+ owner: Class[_],
+ name: String,
+ returnType: Class[_],
+ parameterTypes: Class[_]*): Method = {
+ val result = owner.getMethod(name, parameterTypes: _*)
+ assert(
+ !Modifier.isStatic(result.getModifiers),
+ s"${owner.getName}.$name must be an instance method")
+ assert(
+ result.getReturnType == returnType,
+ s"${owner.getName}.$name returns ${result.getReturnType.getName}, expected ${returnType.getName}")
+ result
+ }
+}