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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
