You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "pingzh (via GitHub)" <gi...@apache.org> on 2026/08/27 17:09:19 UTC
[PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
pingzh opened a new pull request, #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501
## Which issue does this PR close?
Part of #5352. This is the fifth foundational PR and does not close the issue.
Previous PRs:
- https://github.com/apache/datafusion-comet/pull/5473
- https://github.com/apache/datafusion-comet/pull/5476
- https://github.com/apache/datafusion-comet/pull/5481
- https://github.com/apache/datafusion-comet/pull/5491
## Rationale for this change
Previous PRs introduced backend-neutral RSS partition writers, destination-aware shuffle plans, native shuffle execution, and task-owned JVM callbacks.
Integrating Apache Celeborn additionally requires a Celeborn-specific partition pusher and a shuffle manager that preserves existing Celeborn behavior for ordinary Spark shuffles.
## What changes are included in this PR?
- Add `CelebornShufflePartitionPusher`, which:
- Sends complete native shuffle frames through Celeborn’s existing client.
- Captures shuffle, mapper, attempt, and partition metadata.
- Validates frame boundaries, partition identifiers, and transport byte counts.
- Preserves original client exceptions across reflective calls.
- Avoids introducing a compile-time Celeborn dependency.
- Add `CometCelebornShuffleManager`, which:
- Loads the existing Celeborn shuffle manager reflectively.
- Delegates ordinary Spark shuffle registration, writing, reading, cleanup, and shutdown.
- Rejects unsupported Comet shuffle paths until subsequent integration PRs.
- Fails clearly when the optional Celeborn client is unavailable.
- Register the new test suites in Linux and macOS CI workflows.
Map-side lifecycle management, Celeborn shuffle reading, and native-only planning will be addressed in subsequent PRs.
## How are these changes tested?
- Added 11 partition-pusher tests covering frame validation, task metadata, transport accounting, client failures, and worker-thread execution.
- Added 12 shuffle-manager tests covering delegation, reader ranges, lifecycle handling, unsupported shuffle paths, and missing Celeborn dependencies.
- Passed all 35 existing native shuffle tests.
- Passed all 58 tests across the three relevant suites.
- Passed ScalaStyle, Spotless, Apache license-header verification, and CI suite-registration checks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501#issuecomment-5444726905
Merged, thanks @pingzh !
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
Posted by "ziting-openai (via GitHub)" <gi...@apache.org>.
ziting-openai commented on PR #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501#issuecomment-5545306997
Sorry for the noise. My automated review on this PR was posted by a misconfigured AI agent that should have been restricted to OpenAI-owned repositories. Please disregard that automated review. My inline review threads are already resolved. This agent will not review non-OpenAI repositories again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
Posted by "ziting-openai (via GitHub)" <gi...@apache.org>.
ziting-openai commented on code in PR #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501#discussion_r3874367284
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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 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");
+ }
+
+ this.shuffleClient = shuffleClient;
+ this.pushOrMergeData = pushMethod;
+ 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 {
+ accepted =
+ (int)
+ pushOrMergeData.invoke(
Review Comment:
Verified fixed in 6a6f3d27: plaintext CRC/byte accounting now runs exactly once before each raw push, with coverage for multiple partitions, failures, encryption, and legacy clients. This thread can be resolved.
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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 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");
+ }
+
+ this.shuffleClient = shuffleClient;
+ this.pushOrMergeData = pushMethod;
+ 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 {
+ 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 expected = length + CELEBORN_BATCH_HEADER_BYTES;
Review Comment:
Verified fixed in 6a6f3d27: encrypted transport expansion is accepted while undersized responses remain rejected, and the new regression covers combined encryption and integrity accounting. This thread can be resolved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao merged PR #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add Celeborn shuffle manager and partition pusher (5/n) [datafusion-comet]
Posted by "ziting-openai (via GitHub)" <gi...@apache.org>.
ziting-openai commented on code in PR #5501:
URL: https://github.com/apache/datafusion-comet/pull/5501#discussion_r3874152928
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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 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");
+ }
+
+ this.shuffleClient = shuffleClient;
+ this.pushOrMergeData = pushMethod;
+ 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 {
+ accepted =
+ (int)
+ pushOrMergeData.invoke(
Review Comment:
[P2] Record Celeborn's required integrity accounting before raw pushes
With `spark.celeborn.client.shuffle.integrityCheck.enabled=true`, invoking `pushOrMergeData` directly skips the only CRC/byte accumulation path. Celeborn v0.7 explicitly documents that external writers must use `pushDataWithCRC`; that method first calls `computeBatchCRC`, and mapper completion later publishes the accumulated per-partition CRCs and byte counts. This adapter submits the frames without updating either value, so reducer-side integrity validation receives missing/incorrect commit metadata and fails. Please invoke Celeborn's CRC accounting exactly once before each raw push (while preserving the required `skipCompress=true` behavior) and add a regression with integrity checks enabled.
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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 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");
+ }
+
+ this.shuffleClient = shuffleClient;
+ this.pushOrMergeData = pushMethod;
+ 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 {
+ 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 expected = length + CELEBORN_BATCH_HEADER_BYTES;
Review Comment:
[P2] Do not reject successful encrypted Celeborn pushes
With supported Spark I/O encryption enabled (`spark.io.encryption.enabled=true`), Celeborn applies its `CryptoHandler` inside `pushOrMergeData` even when `skipCompress=true`. The encrypted payload is longer than the original Comet frame, and the real client returns `encrypted.length + 16`, not `length + 16`. This exact comparison consequently throws after every encrypted frame has already been submitted, failing/retrying the task and potentially duplicating remote data. Either account for the actual encrypted transport length or reject incompatible encryption before pushing, and cover a length-expanding crypto handler in the test client.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org