jerqi commented on code in PR #293: URL: https://github.com/apache/incubator-uniffle/pull/293#discussion_r1014115486
########## common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java: ########## @@ -0,0 +1,117 @@ +/* + * 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.uniffle.common.segment; + +import java.nio.BufferUnderflowException; +import java.nio.ByteBuffer; +import java.util.List; + +import com.google.common.collect.Lists; +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +import org.apache.uniffle.common.BufferSegment; +import org.apache.uniffle.common.ShuffleDataSegment; +import org.apache.uniffle.common.ShuffleIndexResult; +import org.apache.uniffle.common.exception.RssException; + +public class LocalOrderSegmentSplitter implements SegmentSplitter { + + private Roaring64NavigableMap expectTaskIds; + private int readBufferSize; + + public LocalOrderSegmentSplitter(Roaring64NavigableMap expectTaskIds, int readBufferSize) { + this.expectTaskIds = expectTaskIds; + this.readBufferSize = readBufferSize; + } + + @Override + public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) { + if (shuffleIndexResult == null || shuffleIndexResult.isEmpty()) { + return Lists.newArrayList(); + } + + byte[] indexData = shuffleIndexResult.getIndexData(); + long dataFileLen = shuffleIndexResult.getDataFileLen(); + + ByteBuffer byteBuffer = ByteBuffer.wrap(indexData); + List<BufferSegment> bufferSegments = Lists.newArrayList(); + + List<ShuffleDataSegment> dataFileSegments = Lists.newArrayList(); + int bufferOffset = 0; + long fileOffset = -1; + long totalLen = 0; + + long lastTaskAttemptId = -1; + + /** + * One ShuffleDataSegment should meet following requirements: + * + * 1. taskId in [startMapId, endMapId) taskIds bitmap + * 2. ShuffleDataSegment size should < readBufferSize + * 3. ShuffleDataSegment's blocks should be continuous + * + */ + while (byteBuffer.hasRemaining()) { + try { + long offset = byteBuffer.getLong(); + int length = byteBuffer.getInt(); + int uncompressLength = byteBuffer.getInt(); + long crc = byteBuffer.getLong(); + long blockId = byteBuffer.getLong(); + long taskAttemptId = byteBuffer.getLong(); + + if (lastTaskAttemptId == -1) { + lastTaskAttemptId = taskAttemptId; + } + + // If ShuffleServer is flushing the file at this time, the length in the index file record may be greater + // than the length in the actual data file, and it needs to be returned at this time to avoid EOFException + if (dataFileLen != -1 && totalLen >= dataFileLen) { + break; + } + + if ((taskAttemptId < lastTaskAttemptId && bufferSegments.size() > 0) || bufferOffset >= readBufferSize) { + ShuffleDataSegment sds = new ShuffleDataSegment(fileOffset, bufferOffset, bufferSegments); + dataFileSegments.add(sds); + bufferSegments = Lists.newArrayList(); + bufferOffset = 0; + fileOffset = -1; + } + + if (expectTaskIds.contains(taskAttemptId)) { Review Comment: ExpectTaskId should be a range. The range is between startMapId and endMapId. If a blockId A isn't in this range, there won't be any other blockId after blockId A in this ShuffleDataSegement. Because data is local order in the ShuffleDataSegment. So there won't discontinuous blockIds in a ShuffleDataSegement. The discontinuous blockIds in a ShuffleDataSegment isn't permitted in our system. You can see related logic for details in https://github.com/apache/incubator-uniffle/blob/c69f1735e5bf93852a82de6202e1e5b7c1bd7092/client/src/main/java/org/apache/uniffle/client/impl/ShuffleReadClientImpl.java#L182 -- 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]
