[CELEBORN-627][FLINK] Support split partitions

### What changes were proposed in this pull request? In MapPartiitoin, datas are split into regions. 1. Unlike ReducePartition whose partition split can occur on pushing data to keep MapPartition data ordering, PartitionSplit only be done on the time of sending PushDataHandShake or RegionStart messages (As shown in the following image). That's to say that the partition split only appear at the beginnig of a region but not inner a region. > Notice: if the client side think that it's failed to push HandShake or RegionStart messages. but the worker side can still receive normal HandShake/RegionStart message. After client revive succss, it don't push any messages to old partition, so the worker having the old partition will create a empty file. After committing files, the worker will return empty commitids. That's to say that empty file will be filterd after committing files and ReduceTask will not read any empty files. ![image](https://github.com/apache/incubator-celeborn/assets/96606293/468fd660-afbc-42c1-b111-6643f5c1e944) 2. PushData/RegioinFinish don't care the following cases: - Diskfull - ExceedPartitionSplitThreshold - Worker ShuttingDown so if one of the above three conditions appears, PushData and RegionFinish cant still do as normal. Workers should consider the ShuttingDown case and try best to wait all the regions finished before shutting down. if PushData or RegionFinish failed like network timeout and so on, then MapTask will failed and start another attempte maptask. ![image](https://github.com/apache/incubator-celeborn/assets/96606293/db9f9166-2085-4be1-b09e-cf73b469c55b) 3. how shuffle read supports partition split? ReduceTask should get split paritions by order and open the stream by partition epoc orderly ### Why are the changes needed? PartiitonSplit is not supported by MapPartition from now. There still a risk that a partition file'size is too large to store the file on worker disk. To avoid this risk, this pr introduces partition split in shuffle read and shuffle write. ### Does this PR introduce _any_ user-facing change? NO. ### How was this patch tested? UT and manual TPCDS test Closes #1550 from FMX/CELEBORN-627. Lead-authored-by: zhongqiang.czq <zhongqiang.czq@alibaba-inc.com> Co-authored-by: mingji <fengmingxiao.fmx@alibaba-inc.com> Co-authored-by: Ethan Feng <ethanfeng@apache.org> Signed-off-by: zhongqiang.czq <zhongqiang.czq@alibaba-inc.com>
2023-09-01 19:25:51 +08:00 · 2023-09-01 19:25:51 +08:00 · b66eaff880
commit b66eaff880
parent 28449630f3
26 changed files with 672 additions and 234 deletions
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
@ -24,6 +24,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
 import org.apache.celeborn.common.network.protocol.ReadAddCredit;
 import org.apache.celeborn.common.network.protocol.RequestMessage;
 import org.apache.celeborn.common.network.protocol.TransportableError;
@ -74,13 +75,15 @@ public class RemoteBufferStreamReader extends CreditListener {
            backlogReceived(((BacklogAnnouncement) requestMessage).getBacklog());
          } else if (requestMessage instanceof TransportableError) {
            errorReceived(((TransportableError) requestMessage).getErrorMessage());
+          } else if (requestMessage instanceof BufferStreamEnd) {
+            onStreamEnd((BufferStreamEnd) requestMessage);
          }
        };
  }

  public void open(int initialCredit) {
    try {
-      this.bufferStream =
+      bufferStream =
          client.readBufferedPartition(
              shuffleId, partitionId, subPartitionIndexStart, subPartitionIndexEnd);
      bufferStream.open(
@ -95,7 +98,8 @@ public class RemoteBufferStreamReader extends CreditListener {
  public void close() {
    // need set closed first before remove Handler
    closed = true;
-    if (this.bufferStream != null) {
+    if (bufferStream != null) {
+      logger.debug("Close bufferStream currentStreamId:{}", bufferStream.getStreamId());
      bufferStream.close();
    } else {
      logger.warn(
@ -111,7 +115,7 @@ public class RemoteBufferStreamReader extends CreditListener {

  public void notifyAvailableCredits(int numCredits) {
    if (!closed) {
-      ReadAddCredit addCredit = new ReadAddCredit(this.bufferStream.getStreamId(), numCredits);
+      ReadAddCredit addCredit = new ReadAddCredit(bufferStream.getStreamId(), numCredits);
      bufferStream.addCredit(addCredit);
    }
  }
@ -146,4 +150,10 @@ public class RemoteBufferStreamReader extends CreditListener {
        readData.getFlinkBuffer().readableBytes());
    dataListener.accept(readData.getFlinkBuffer());
  }
+
+  public void onStreamEnd(BufferStreamEnd streamEnd) {
+    long streamId = streamEnd.getStreamId();
+    logger.debug("Buffer stream reader get stream end for {}", streamId);
+    bufferStream.moveToNextPartitionIfPossible(streamId);
+  }
 }
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
@ -79,7 +79,9 @@ public class RemoteShuffleOutputGate {
  private int lifecycleManagerPort;
  private long lifecycleManagerTimestamp;
  private UserIdentifier userIdentifier;
-  private boolean isFirstHandShake = true;
+  private boolean isRegisterShuffle = false;
+  private int maxReviveTimes;
+  private boolean hasSentHandshake = false;

  /**
   * @param shuffleDesc Describes shuffle meta and shuffle worker address.
@ -114,6 +116,7 @@ public class RemoteShuffleOutputGate {
    this.lifecycleManagerTimestamp =
        shuffleDesc.getShuffleResource().getLifecycleManagerTimestamp();
    this.flinkShuffleClient = getShuffleClient();
+    this.maxReviveTimes = celebornConf.clientPushMaxReviveTimes();
  }

  /** Initialize transportation gate. */
@ -144,31 +147,10 @@ public class RemoteShuffleOutputGate {
   * @param isBroadcast Whether it's a broadcast region.
   */
  public void regionStart(boolean isBroadcast) {
-    Optional<PartitionLocation> newPartitionLoc;
    try {
-      if (isFirstHandShake) {
-        handshake(true);
-        isFirstHandShake = false;
-        LOG.debug(
-            "shuffleId: {}, location: {}, send firstHandShake: {}, isBroadcast: {}",
-            shuffleId,
-            partitionLocation.getUniqueId(),
-            true,
-            isBroadcast);
-      }
-
-      newPartitionLoc =
-          flinkShuffleClient.regionStart(
-              shuffleId, mapId, attemptId, partitionLocation, currentRegionIndex, isBroadcast);
-      // revived
-      if (newPartitionLoc.isPresent()) {
-        partitionLocation = newPartitionLoc.get();
-        // send handshake again
-        handshake(false);
-        // send regionstart again
-        flinkShuffleClient.regionStart(
-            shuffleId, mapId, attemptId, newPartitionLoc.get(), currentRegionIndex, isBroadcast);
-      }
+      registerShuffle();
+      handshake();
+      regionStartWithRevive(isBroadcast);
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
@ -240,18 +222,86 @@ public class RemoteShuffleOutputGate {
    }
  }

-  public void handshake(boolean isFirstHandShake) throws IOException {
-    if (isFirstHandShake) {
+  public void registerShuffle() throws IOException {
+    if (!isRegisterShuffle) {
      partitionLocation =
          flinkShuffleClient.registerMapPartitionTask(
              shuffleId, numMappers, mapId, attemptId, partitionId);
      Utils.checkNotNull(partitionLocation);

      currentRegionIndex = 0;
+      isRegisterShuffle = true;
    }
+  }
+
+  public void regionStartWithRevive(boolean isBroadcast) {
    try {
-      flinkShuffleClient.pushDataHandShake(
-          shuffleId, mapId, attemptId, numSubs, bufferSize, partitionLocation);
+      int remainingReviveTimes = maxReviveTimes;
+      boolean hasSentRegionStart = false;
+      while (remainingReviveTimes-- > 0 && !hasSentRegionStart) {
+        Optional<PartitionLocation> revivePartition =
+            flinkShuffleClient.regionStart(
+                shuffleId, mapId, attemptId, partitionLocation, currentRegionIndex, isBroadcast);
+        if (revivePartition.isPresent()) {
+          LOG.info(
+              "Revive at regionStart, currentTimes:{}, totalTimes:{} for shuffleId:{}, mapId:{}, attempId:{}, currentRegionIndex:{}, isBroadcast:{}, newPartition:{}, oldPartition:{}",
+              remainingReviveTimes,
+              maxReviveTimes,
+              shuffleId,
+              mapId,
+              attemptId,
+              currentRegionIndex,
+              isBroadcast,
+              revivePartition,
+              partitionLocation);
+          partitionLocation = revivePartition.get();
+          hasSentRegionStart = false;
+          // For every revive partition, handshake should be sent firstly
+          hasSentHandshake = false;
+          handshake();
+        } else {
+          hasSentRegionStart = true;
+        }
+      }
+      if (remainingReviveTimes == 0 && !hasSentRegionStart) {
+        throw new RuntimeException(
+            "After retry " + maxReviveTimes + " times, still failed to send regionStart");
+      }
+    } catch (IOException e) {
+      Utils.rethrowAsRuntimeException(e);
+    }
+  }
+
+  public void handshake() {
+    try {
+      int remainingReviveTimes = maxReviveTimes;
+      while (remainingReviveTimes-- > 0 && !hasSentHandshake) {
+        Optional<PartitionLocation> revivePartition =
+            flinkShuffleClient.pushDataHandShake(
+                shuffleId, mapId, attemptId, numSubs, bufferSize, partitionLocation);
+        // if remainingReviveTimes == 0 and revivePartition.isPresent(), there is no need to send
+        // handshake again
+        if (revivePartition.isPresent() && remainingReviveTimes > 0) {
+          LOG.info(
+              "Revive at handshake, currentTimes:{}, totalTimes:{} for shuffleId:{}, mapId:{}, attempId:{}, currentRegionIndex:{}, newPartition:{}, oldPartition:{}",
+              remainingReviveTimes,
+              maxReviveTimes,
+              shuffleId,
+              mapId,
+              attemptId,
+              currentRegionIndex,
+              revivePartition,
+              partitionLocation);
+          partitionLocation = revivePartition.get();
+          hasSentHandshake = false;
+        } else {
+          hasSentHandshake = true;
+        }
+      }
+      if (remainingReviveTimes == 0 && !hasSentHandshake) {
+        throw new RuntimeException(
+            "After retry " + maxReviveTimes + " times, still failed to send handshake");
+      }
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java
@ -89,6 +89,10 @@ public class MessageDecoderExt {
      case HEARTBEAT:
        return new Heartbeat();

+      case BUFFER_STREAM_END:
+        streamId = in.readLong();
+        return new BufferStreamEnd(streamId);
+
      default:
        throw new IllegalArgumentException("Unexpected message type: " + type);
    }
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
@ -26,6 +26,7 @@ import org.slf4j.LoggerFactory;

 import org.apache.celeborn.common.network.client.TransportClient;
 import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
 import org.apache.celeborn.common.network.protocol.RequestMessage;
 import org.apache.celeborn.common.network.protocol.TransportableError;
 import org.apache.celeborn.common.network.server.BaseMessageHandler;
@ -86,6 +87,11 @@ public class ReadClientHandler extends BaseMessageHandler {
            transportableError.getErrorMessage());
        processMessageInternal(streamId, transportableError);
        break;
+      case BUFFER_STREAM_END:
+        BufferStreamEnd streamEnd = (BufferStreamEnd) msg;
+        logger.debug("Received streamend for {}", streamEnd.getStreamId());
+        processMessageInternal(streamEnd.getStreamId(), streamEnd);
+        break;
      case ONE_WAY_MESSAGE:
        // ignore it.
        break;
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/CelebornBufferStream.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/CelebornBufferStream.java
@ -19,6 +19,7 @@ package org.apache.celeborn.plugin.flink.readclient;

 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import java.util.function.Supplier;

@ -26,43 +27,44 @@ import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

-import org.apache.celeborn.common.CelebornConf;
 import org.apache.celeborn.common.network.client.RpcResponseCallback;
 import org.apache.celeborn.common.network.client.TransportClient;
 import org.apache.celeborn.common.network.protocol.*;
 import org.apache.celeborn.common.network.util.NettyUtils;
+import org.apache.celeborn.common.protocol.MessageType;
 import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbOpenStream;
+import org.apache.celeborn.common.protocol.PbStreamHandler;
 import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory;

 public class CelebornBufferStream {
-
  private static Logger logger = LoggerFactory.getLogger(CelebornBufferStream.class);
-  private CelebornConf conf;
  private FlinkTransportClientFactory clientFactory;
  private String shuffleKey;
  private PartitionLocation[] locations;
  private int subIndexStart;
  private int subIndexEnd;
  private TransportClient client;
-  private int currentLocationIndex = 0;
+  private AtomicInteger currentLocationIndex = new AtomicInteger(0);
  private long streamId = 0;
  private FlinkShuffleClientImpl mapShuffleClient;
  private boolean isClosed;
  private boolean isOpenSuccess;
  private Object lock = new Object();
+  private Supplier<ByteBuf> bufferSupplier;
+  private int initialCredit;
+  private Consumer<RequestMessage> messageConsumer;

  public CelebornBufferStream() {}

  public CelebornBufferStream(
      FlinkShuffleClientImpl mapShuffleClient,
-      CelebornConf conf,
      FlinkTransportClientFactory dataClientFactory,
      String shuffleKey,
      PartitionLocation[] locations,
      int subIndexStart,
      int subIndexEnd) {
    this.mapShuffleClient = mapShuffleClient;
-    this.conf = conf;
    this.clientFactory = dataClientFactory;
    this.shuffleKey = shuffleKey;
    this.locations = locations;
@ -71,56 +73,13 @@ public class CelebornBufferStream {
  }

  public void open(
-      Supplier<ByteBuf> supplier, int initialCredit, Consumer<RequestMessage> messageConsumer)
-      throws IOException, InterruptedException {
-    this.client =
-        clientFactory.createClientWithRetry(
-            locations[currentLocationIndex].getHost(),
-            locations[currentLocationIndex].getFetchPort());
-    String fileName = locations[currentLocationIndex].getFileName();
-    OpenStreamWithCredit openBufferStream =
-        new OpenStreamWithCredit(shuffleKey, fileName, subIndexStart, subIndexEnd, initialCredit);
-    client.sendRpc(
-        openBufferStream.toByteBuffer(),
-        new RpcResponseCallback() {
-
-          @Override
-          public void onSuccess(ByteBuffer response) {
-            StreamHandle streamHandle = (StreamHandle) Message.decode(response);
-            CelebornBufferStream.this.streamId = streamHandle.streamId;
-            synchronized (lock) {
-              if (!isClosed) {
-                clientFactory.registerSupplier(CelebornBufferStream.this.streamId, supplier);
-                mapShuffleClient
-                    .getReadClientHandler()
-                    .registerHandler(streamId, messageConsumer, client);
-                isOpenSuccess = true;
-                logger.debug(
-                    "open stream success from remote:{}, stream id:{}, fileName: {}",
-                    client.getSocketAddress(),
-                    streamId,
-                    fileName);
-              } else {
-                logger.debug(
-                    "open stream success from remote:{}, but stream reader is already closed, stream id:{}, fileName: {}",
-                    client.getSocketAddress(),
-                    streamId,
-                    fileName);
-                closeStream();
-              }
-            }
-          }
-
-          @Override
-          public void onFailure(Throwable e) {
-            logger.error(
-                "Open file {} stream for {} error from {}",
-                fileName,
-                shuffleKey,
-                NettyUtils.getRemoteAddress(client.getChannel()));
-            messageConsumer.accept(new TransportableError(streamId, e));
-          }
-        });
+      Supplier<ByteBuf> bufferSupplier,
+      int initialCredit,
+      Consumer<RequestMessage> messageConsumer) {
+    this.bufferSupplier = bufferSupplier;
+    this.initialCredit = initialCredit;
+    this.messageConsumer = messageConsumer;
+    moveToNextPartitionIfPossible(0);
  }

  public void addCredit(ReadAddCredit addCredit) {
@ -150,7 +109,6 @@ public class CelebornBufferStream {

  public static CelebornBufferStream create(
      FlinkShuffleClientImpl client,
-      CelebornConf conf,
      FlinkTransportClientFactory dataClientFactory,
      String shuffleKey,
      PartitionLocation[] locations,
@ -160,30 +118,135 @@ public class CelebornBufferStream {
      return empty();
    } else {
      return new CelebornBufferStream(
-          client, conf, dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd);
+          client, dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd);
    }
  }

  private static final CelebornBufferStream EMPTY_CELEBORN_BUFFER_STREAM =
      new CelebornBufferStream();

-  private void closeStream() {
+  private void closeStream(long streamId) {
    if (client != null && client.isActive()) {
      client.getChannel().writeAndFlush(new BufferStreamEnd(streamId));
    }
  }

+  private void cleanStream(long streamId) {
+    if (isOpenSuccess) {
+      mapShuffleClient.getReadClientHandler().removeHandler(streamId);
+      clientFactory.unregisterSupplier(streamId);
+      closeStream(streamId);
+      isOpenSuccess = false;
+    }
+  }
+
  public void close() {
    synchronized (lock) {
-      if (isOpenSuccess) {
-        mapShuffleClient.getReadClientHandler().removeHandler(getStreamId());
-        clientFactory.unregisterSupplier(this.getStreamId());
-        closeStream();
-      }
+      cleanStream(streamId);
      isClosed = true;
    }
  }

+  public void moveToNextPartitionIfPossible(long endedStreamId) {
+    logger.debug(
+        "MoveToNextPartitionIfPossible in this:{},  endedStreamId: {}, currentLocationIndex: {}, currentSteamId:{}, locationsLength:{}",
+        this,
+        endedStreamId,
+        currentLocationIndex.get(),
+        streamId,
+        locations.length);
+    if (currentLocationIndex.get() > 0) {
+      logger.debug("Get end streamId {}", endedStreamId);
+      cleanStream(endedStreamId);
+    }
+    if (currentLocationIndex.get() < locations.length) {
+      try {
+        openStreamInternal();
+        logger.debug(
+            "MoveToNextPartitionIfPossible after openStream this:{},  endedStreamId: {}, currentLocationIndex: {}, currentSteamId:{}, locationsLength:{}",
+            this,
+            endedStreamId,
+            currentLocationIndex.get(),
+            streamId,
+            locations.length);
+      } catch (Exception e) {
+        logger.warn("Failed to open stream and report to flink framework. ", e);
+        messageConsumer.accept(new TransportableError(0L, e));
+      }
+    }
+  }
+
+  private void openStreamInternal() throws IOException, InterruptedException {
+    this.client =
+        clientFactory.createClientWithRetry(
+            locations[currentLocationIndex.get()].getHost(),
+            locations[currentLocationIndex.get()].getFetchPort());
+    String fileName = locations[currentLocationIndex.getAndIncrement()].getFileName();
+    TransportMessage openStream =
+        new TransportMessage(
+            MessageType.OPEN_STREAM,
+            PbOpenStream.newBuilder()
+                .setShuffleKey(shuffleKey)
+                .setFileName(fileName)
+                .setStartIndex(subIndexStart)
+                .setEndIndex(subIndexEnd)
+                .setInitialCredit(initialCredit)
+                .build()
+                .toByteArray());
+    client.sendRpc(
+        openStream.toByteBuffer(),
+        new RpcResponseCallback() {
+
+          @Override
+          public void onSuccess(ByteBuffer response) {
+            try {
+              PbStreamHandler pbStreamHandler =
+                  TransportMessage.fromByteBuffer(response).getParsedPayload();
+              CelebornBufferStream.this.streamId = pbStreamHandler.getStreamId();
+              synchronized (lock) {
+                if (!isClosed) {
+                  clientFactory.registerSupplier(
+                      CelebornBufferStream.this.streamId, bufferSupplier);
+                  mapShuffleClient
+                      .getReadClientHandler()
+                      .registerHandler(streamId, messageConsumer, client);
+                  isOpenSuccess = true;
+                  logger.debug(
+                      "open stream success from remote:{}, stream id:{}, fileName: {}",
+                      client.getSocketAddress(),
+                      streamId,
+                      fileName);
+                } else {
+                  logger.debug(
+                      "open stream success from remote:{}, but stream reader is already closed, stream id:{}, fileName: {}",
+                      client.getSocketAddress(),
+                      streamId,
+                      fileName);
+                  closeStream(streamId);
+                }
+              }
+            } catch (Exception e) {
+              logger.error(
+                  "Open file {} stream for {} error from {}",
+                  fileName,
+                  shuffleKey,
+                  NettyUtils.getRemoteAddress(client.getChannel()));
+              messageConsumer.accept(new TransportableError(streamId, e));
+            }
+          }
+
+          @Override
+          public void onFailure(Throwable e) {
+            logger.error(
+                "Open file {} stream for {} error from {}",
+                fileName,
+                shuffleKey,
+                NettyUtils.getRemoteAddress(client.getChannel()));
+            messageConsumer.accept(new TransportableError(streamId, e));
+          }
+        });
+  }
+
  public TransportClient getClient() {
    return client;
  }
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
@ -148,12 +148,19 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
      logger.error("Shuffle data is empty for shuffle {} partitionId {}.", shuffleId, partitionId);
      throw new PartitionUnRetryAbleException(partitionId + " may be lost.");
    } else {
+      PartitionLocation[] partitionLocations =
+          fileGroups.partitionGroups.get(partitionId).toArray(new PartitionLocation[0]);
+      Arrays.sort(partitionLocations, Comparator.comparingInt(PartitionLocation::getEpoch));
+      logger.debug(
+          "readBufferedPartition shuffleKey:{} partitionid:{} partitionLocation:{}",
+          shuffleKey,
+          partitionId,
+          partitionLocations);
      return CelebornBufferStream.create(
          this,
-          conf,
          flinkTransportClientFactory,
          shuffleKey,
-          fileGroups.partitionGroups.get(partitionId).toArray(new PartitionLocation[0]),
+          partitionLocations,
          subPartitionIndexStart,
          subPartitionIndexEnd);
    }
@ -305,7 +312,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
    return currentClient.get(mapKey);
  }

-  public void pushDataHandShake(
+  public Optional<PartitionLocation> pushDataHandShake(
      int shuffleId,
      int mapId,
      int attemptId,
@ -315,12 +322,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
      throws IOException {
    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
    final PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
-    sendMessageInternal(
-        shuffleId,
-        mapId,
-        attemptId,
-        location,
-        pushState,
+    return retrySendMessage(
        () -> {
          String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
          logger.info(
@ -338,8 +340,20 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
                  attemptId,
                  numPartitions,
                  bufferSize);
-          client.sendRpcSync(handShake.toByteBuffer(), conf.pushDataTimeoutMs());
-          return null;
+          ByteBuffer pushDataHandShakeResponse;
+          try {
+            pushDataHandShakeResponse =
+                client.sendRpcSync(handShake.toByteBuffer(), conf.pushDataTimeoutMs());
+          } catch (IOException e) {
+            // ioexeption revive
+            return revive(shuffleId, mapId, attemptId, location);
+          }
+          if (pushDataHandShakeResponse.hasRemaining()
+              && pushDataHandShakeResponse.get() == StatusCode.HARD_SPLIT.getValue()) {
+            // if split then revive
+            return revive(shuffleId, mapId, attemptId, location);
+          }
+          return Optional.empty();
        });
  }

@ -353,12 +367,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
      throws IOException {
    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
    final PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
-    return sendMessageInternal(
-        shuffleId,
-        mapId,
-        attemptId,
-        location,
-        pushState,
+    return retrySendMessage(
        () -> {
          String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
          logger.info(
@ -377,61 +386,69 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
                  attemptId,
                  currentRegionIdx,
                  isBroadcast);
-          ByteBuffer regionStartResponse =
-              client.sendRpcSync(regionStart.toByteBuffer(), conf.pushDataTimeoutMs());
+          ByteBuffer regionStartResponse;
+          try {
+            regionStartResponse =
+                client.sendRpcSync(regionStart.toByteBuffer(), conf.pushDataTimeoutMs());
+          } catch (IOException e) {
+            // ioexeption revive
+            return revive(shuffleId, mapId, attemptId, location);
+          }
+
          if (regionStartResponse.hasRemaining()
              && regionStartResponse.get() == StatusCode.HARD_SPLIT.getValue()) {
            // if split then revive
-            Set<Integer> mapIds = new HashSet<>();
-            mapIds.add(mapId);
-            List<ReviveRequest> requests = new ArrayList<>();
-            ReviveRequest req =
-                new ReviveRequest(
-                    shuffleId,
-                    mapId,
-                    attemptId,
-                    location.getId(),
-                    location.getEpoch(),
-                    location,
-                    StatusCode.HARD_SPLIT);
-            requests.add(req);
-            PbChangeLocationResponse response =
-                lifecycleManagerRef.askSync(
-                    ControlMessages.Revive$.MODULE$.apply(shuffleId, mapIds, requests),
-                    conf.clientRpcRequestPartitionLocationRpcAskTimeout(),
-                    ClassTag$.MODULE$.apply(PbChangeLocationResponse.class));
-            // per partitionKey only serve single PartitionLocation in Client Cache.
-            PbChangeLocationPartitionInfo partitionInfo = response.getPartitionInfo(0);
-            StatusCode respStatus = Utils.toStatusCode(partitionInfo.getStatus());
-            if (StatusCode.SUCCESS.equals(respStatus)) {
-              return Optional.of(
-                  PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition()));
-            } else {
-              // throw exception
-              logger.error(
-                  "Exception raised while reviving for shuffle {} map {} attemptId {} partition {} epoch {}.",
-                  shuffleId,
-                  mapId,
-                  attemptId,
-                  location.getId(),
-                  location.getEpoch());
-              throw new CelebornIOException("RegionStart revive failed");
-            }
+            return revive(shuffleId, mapId, attemptId, location);
          }
          return Optional.empty();
        });
  }

+  public Optional<PartitionLocation> revive(
+      int shuffleId, int mapId, int attemptId, PartitionLocation location)
+      throws CelebornIOException {
+    Set<Integer> mapIds = new HashSet<>();
+    mapIds.add(mapId);
+    List<ReviveRequest> requests = new ArrayList<>();
+    ReviveRequest req =
+        new ReviveRequest(
+            shuffleId,
+            mapId,
+            attemptId,
+            location.getId(),
+            location.getEpoch(),
+            location,
+            StatusCode.HARD_SPLIT);
+    requests.add(req);
+    PbChangeLocationResponse response =
+        lifecycleManagerRef.askSync(
+            ControlMessages.Revive$.MODULE$.apply(shuffleId, mapIds, requests),
+            conf.clientRpcRequestPartitionLocationRpcAskTimeout(),
+            ClassTag$.MODULE$.apply(PbChangeLocationResponse.class));
+    // per partitionKey only serve single PartitionLocation in Client Cache.
+    PbChangeLocationPartitionInfo partitionInfo = response.getPartitionInfo(0);
+    StatusCode respStatus = Utils.toStatusCode(partitionInfo.getStatus());
+    if (StatusCode.SUCCESS.equals(respStatus)) {
+      logger.debug("revive new partition:{}", partitionInfo.getPartition());
+      return Optional.of(PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition()));
+    } else {
+      // throw exception
+      logger.error(
+          "Exception raised while reviving for shuffle {} map {} attemptId {} partition {} epoch {}.",
+          shuffleId,
+          mapId,
+          attemptId,
+          location.getId(),
+          location.getEpoch());
+      throw new CelebornIOException("RegionStart revive failed");
+    }
+  }
+
  public void regionFinish(int shuffleId, int mapId, int attemptId, PartitionLocation location)
      throws IOException {
    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
    final PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
-    sendMessageInternal(
-        shuffleId,
-        mapId,
-        attemptId,
-        location,
-        pushState,
+    retrySendMessage(
        () -> {
          final String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
          logger.info(
@ -449,31 +466,6 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
        });
  }

-  private <R> R sendMessageInternal(
-      int shuffleId,
-      int mapId,
-      int attemptId,
-      PartitionLocation location,
-      PushState pushState,
-      ThrowingExceptionSupplier<R, Exception> supplier)
-      throws IOException {
-    int batchId = 0;
-    try {
-      // mapKey
-      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
-      pushState = getPushState(mapKey);
-
-      // add inFlight requests
-      batchId = pushState.nextBatchId();
-      pushState.addBatch(batchId, location.hostAndPushPort());
-      return retrySendMessage(supplier);
-    } finally {
-      if (pushState != null) {
-        pushState.removeBatch(batchId, location.hostAndPushPort());
-      }
-    }
-  }
-
  @FunctionalInterface
  interface ThrowingExceptionSupplier<R, E extends Exception> {
    R get() throws E;
--- a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java
@ -56,11 +56,11 @@ public class RemoteShuffleOutputGateSuiteJ {
            1, 0, "localhost", 123, 245, 789, 238, PartitionLocation.Mode.PRIMARY);
    when(shuffleClient.registerMapPartitionTask(anyInt(), anyInt(), anyInt(), anyInt(), anyInt()))
        .thenAnswer(t -> partitionLocation);
-    doNothing()
-        .when(remoteShuffleOutputGate.flinkShuffleClient)
-        .pushDataHandShake(anyInt(), anyInt(), anyInt(), anyInt(), anyInt(), any());
+    when(remoteShuffleOutputGate.flinkShuffleClient.pushDataHandShake(
+            anyInt(), anyInt(), anyInt(), anyInt(), anyInt(), any()))
+        .thenAnswer(t -> Optional.empty());

-    remoteShuffleOutputGate.handshake(true);
+    remoteShuffleOutputGate.handshake();

    when(remoteShuffleOutputGate.flinkShuffleClient.regionStart(
            anyInt(), anyInt(), anyInt(), any(), anyInt(), anyBoolean()))
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@ -499,7 +499,8 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
        false)
      return
    }
-
+    logDebug(
+      s"[handleRevive] shuffle $shuffleId, $mapIds, $partitionIds, $oldEpochs, $oldPartitions, $causes")
    if (commitManager.isStageEnd(shuffleId)) {
      logError(s"[handleRevive] shuffle $shuffleId stage ended!")
      contextWrapper.reply(
@ -662,7 +663,10 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
            getPartitionType(shuffleId),
            rangeReadFilter,
            userIdentifier,
-            conf.pushDataTimeoutMs))
+            conf.pushDataTimeoutMs,
+            if (getPartitionType(shuffleId) == PartitionType.MAP)
+              conf.clientShuffleMapPartitionSplitEnabled
+            else true))
        if (res.status.equals(StatusCode.SUCCESS)) {
          logDebug(s"Successfully allocated " +
            s"partitions buffer for shuffleId $shuffleId" +
--- a/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
@ -48,20 +48,25 @@ public class FileInfo {
  private int numSubpartitions;

  private volatile long bytesFlushed;
+  // whether to split is decided by client side.
+  // now it's just used for mappartition to compatible with old client which can't support split
+  private boolean partitionSplitEnabled;

  public FileInfo(String filePath, List<Long> chunkOffsets, UserIdentifier userIdentifier) {
-    this(filePath, chunkOffsets, userIdentifier, PartitionType.REDUCE);
+    this(filePath, chunkOffsets, userIdentifier, PartitionType.REDUCE, true);
  }

  public FileInfo(
      String filePath,
      List<Long> chunkOffsets,
      UserIdentifier userIdentifier,
-      PartitionType partitionType) {
+      PartitionType partitionType,
+      boolean partitionSplitEnabled) {
    this.filePath = filePath;
    this.chunkOffsets = chunkOffsets;
    this.userIdentifier = userIdentifier;
    this.partitionType = partitionType;
+    this.partitionSplitEnabled = partitionSplitEnabled;
  }

  public FileInfo(
@ -71,7 +76,8 @@ public class FileInfo {
      PartitionType partitionType,
      int bufferSize,
      int numSubpartitions,
-      long bytesFlushed) {
+      long bytesFlushed,
+      boolean partitionSplitEnabled) {
    this.filePath = filePath;
    this.chunkOffsets = chunkOffsets;
    this.userIdentifier = userIdentifier;
@ -79,10 +85,24 @@ public class FileInfo {
    this.bufferSize = bufferSize;
    this.numSubpartitions = numSubpartitions;
    this.bytesFlushed = bytesFlushed;
+    this.partitionSplitEnabled = partitionSplitEnabled;
  }

  public FileInfo(String filePath, UserIdentifier userIdentifier, PartitionType partitionType) {
-    this(filePath, new ArrayList(Arrays.asList(0L)), userIdentifier, partitionType);
+    this(filePath, new ArrayList(Arrays.asList(0L)), userIdentifier, partitionType, true);
+  }
+
+  public FileInfo(
+      String filePath,
+      UserIdentifier userIdentifier,
+      PartitionType partitionType,
+      boolean partitionSplitEnabled) {
+    this(
+        filePath,
+        new ArrayList(Arrays.asList(0L)),
+        userIdentifier,
+        partitionType,
+        partitionSplitEnabled);
  }

  @VisibleForTesting
@ -91,7 +111,8 @@ public class FileInfo {
        file.getAbsolutePath(),
        new ArrayList(Arrays.asList(0L)),
        userIdentifier,
-        PartitionType.REDUCE);
+        PartitionType.REDUCE,
+        true);
  }

  public synchronized void addChunkOffset(long bytesFlushed) {
@ -236,4 +257,12 @@ public class FileInfo {
  public long getBytesFlushed() {
    return bytesFlushed;
  }
+
+  public boolean isPartitionSplitEnabled() {
+    return partitionSplitEnabled;
+  }
+
+  public void setPartitionSplitEnabled(boolean partitionSplitEnabled) {
+    this.partitionSplitEnabled = partitionSplitEnabled;
+  }
 }
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@ -349,6 +349,7 @@ message PbReserveSlots {
  bool rangeReadFilter = 8;
  PbUserIdentifier userIdentifier = 9;
  int64 pushDataTimeout = 10;
+  bool partitionSplitEnabled = 11;
 }

 message PbReserveSlotsResponse {
@ -431,6 +432,7 @@ message PbFileInfo {
  int32 bufferSize = 5;
  int32 numSubpartitions = 6;
  int64 bytesFlushed = 7;
+  bool partitionSplitEnabled = 8;
 }

 message PbFileInfoMap {
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@ -1028,6 +1028,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
  def clientFlinkResultPartitionSupportFloatingBuffer: Boolean =
    get(CLIENT_RESULT_PARTITION_SUPPORT_FLOATING_BUFFER)
  def clientFlinkDataCompressionEnabled: Boolean = get(CLIENT_DATA_COMPRESSION_ENABLED)
+  def clientShuffleMapPartitionSplitEnabled = get(CLIENT_SHUFFLE_MAPPARTITION_SPLIT_ENABLED)
 }

 object CelebornConf extends Logging {
@ -3816,4 +3817,13 @@ object CelebornConf extends Logging {
      .doc("Threads count for read local shuffle file.")
      .intConf
      .createWithDefault(4)
+
+  val CLIENT_SHUFFLE_MAPPARTITION_SPLIT_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.client.shuffle.mapPartition.split.enabled")
+      .categories("client")
+      .doc(
+        "whether to enable shuffle partition split. Currently, this only applies to MapPartition.")
+      .version("0.3.1")
+      .booleanConf
+      .createWithDefault(false)
 }
--- a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@ -383,7 +383,8 @@ object ControlMessages extends Logging {
      partitionType: PartitionType,
      rangeReadFilter: Boolean,
      userIdentifier: UserIdentifier,
-      pushDataTimeout: Long)
+      pushDataTimeout: Long,
+      partitionSplitEnabled: Boolean = false)
    extends WorkerMessage

  case class ReserveSlotsResponse(
@ -694,7 +695,8 @@ object ControlMessages extends Logging {
          partType,
          rangeReadFilter,
          userIdentifier,
-          pushDataTimeout) =>
+          pushDataTimeout,
+          partitionSplitEnabled) =>
      val payload = PbReserveSlots.newBuilder()
        .setApplicationId(applicationId)
        .setShuffleId(shuffleId)
@ -708,6 +710,7 @@ object ControlMessages extends Logging {
        .setRangeReadFilter(rangeReadFilter)
        .setUserIdentifier(PbSerDeUtils.toPbUserIdentifier(userIdentifier))
        .setPushDataTimeout(pushDataTimeout)
+        .setPartitionSplitEnabled(partitionSplitEnabled)
        .build().toByteArray
      new TransportMessage(MessageType.RESERVE_SLOTS, payload)

@ -1002,7 +1005,8 @@ object ControlMessages extends Logging {
          Utils.toPartitionType(pbReserveSlots.getPartitionType),
          pbReserveSlots.getRangeReadFilter,
          userIdentifier,
-          pbReserveSlots.getPushDataTimeout)
+          pbReserveSlots.getPushDataTimeout,
+          pbReserveSlots.getPartitionSplitEnabled)

      case RESERVE_SLOTS_RESPONSE_VALUE =>
        val pbReserveSlotsResponse = PbReserveSlotsResponse.parseFrom(message.getPayload)
--- a/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
@ -93,7 +93,8 @@ object PbSerDeUtils {
      Utils.toPartitionType(pbFileInfo.getPartitionType),
      pbFileInfo.getBufferSize,
      pbFileInfo.getNumSubpartitions,
-      pbFileInfo.getBytesFlushed)
+      pbFileInfo.getBytesFlushed,
+      pbFileInfo.getPartitionSplitEnabled)

  def toPbFileInfo(fileInfo: FileInfo): PbFileInfo =
    PbFileInfo.newBuilder
@ -104,6 +105,7 @@ object PbSerDeUtils {
      .setBufferSize(fileInfo.getBufferSize)
      .setNumSubpartitions(fileInfo.getNumSubpartitions)
      .setBytesFlushed(fileInfo.getFileLength)
+      .setPartitionSplitEnabled(fileInfo.isPartitionSplitEnabled)
      .build

  @throws[InvalidProtocolBufferException]
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@ -88,6 +88,7 @@ license: |
 | celeborn.client.shuffle.compression.zstd.level | 1 | Compression level for Zstd compression codec, its value should be an integer between -5 and 22. Increasing the compression level will result in better compression at the expense of more CPU and memory. | 0.3.0 | 
 | celeborn.client.shuffle.expired.checkInterval | 60s | Interval for client to check expired shuffles. | 0.3.0 | 
 | celeborn.client.shuffle.manager.port | 0 | Port used by the LifecycleManager on the Driver. | 0.3.0 | 
+| celeborn.client.shuffle.mapPartition.split.enabled | false | whether to enable shuffle partition split. Currently, this only applies to MapPartition. | 0.3.1 | 
 | celeborn.client.shuffle.partition.type | REDUCE | Type of shuffle's partition. | 0.3.0 | 
 | celeborn.client.shuffle.partitionSplit.mode | SOFT | soft: the shuffle file size might be larger than split threshold. hard: the shuffle file size will be limited to split threshold. | 0.3.0 | 
 | celeborn.client.shuffle.partitionSplit.threshold | 1G | Shuffle file size threshold, if file size exceeds this, trigger split. | 0.3.0 | 
--- a/tests/flink-it/src/test/java/org/apache/celeborn/tests/flink/SplitHelper.java
+++ b/tests/flink-it/src/test/java/org/apache/celeborn/tests/flink/SplitHelper.java
@ -0,0 +1,68 @@
+/*
+ * 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.celeborn.tests.flink;
+
+import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.flink.api.common.functions.FlatMapFunction;
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.sink.SinkFunction;
+import org.apache.flink.util.Collector;
+import org.junit.Assert;
+
+public class SplitHelper {
+  private static final int NUM_WORDS = 10000;
+
+  private static final Long WORD_COUNT = 200L;
+
+  public static void runSplitRead(StreamExecutionEnvironment env) throws Exception {
+    DataStream<Tuple2<String, Long>> words =
+        env.fromSequence(0, NUM_WORDS)
+            .map(
+                new MapFunction<Long, String>() {
+                  @Override
+                  public String map(Long index) throws Exception {
+                    return index + "_" + RandomStringUtils.randomAlphabetic(10);
+                  }
+                })
+            .flatMap(
+                new FlatMapFunction<String, Tuple2<String, Long>>() {
+                  @Override
+                  public void flatMap(String s, Collector<Tuple2<String, Long>> collector)
+                      throws Exception {
+                    for (int i = 0; i < WORD_COUNT; ++i) {
+                      collector.collect(new Tuple2<>(s, 1L));
+                    }
+                  }
+                });
+    words
+        .keyBy(value -> value.f0)
+        .sum(1)
+        .map((MapFunction<Tuple2<String, Long>, Long>) wordCount -> wordCount.f1)
+        .addSink(
+            new SinkFunction<Long>() {
+              @Override
+              public void invoke(Long value, Context context) throws Exception {
+                Assert.assertEquals(value, WORD_COUNT);
+                //                      Thread.sleep(30 * 1000);
+              }
+            });
+  }
+}
--- a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/SplitTest.scala
+++ b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/SplitTest.scala
@ -0,0 +1,78 @@
+/*
+ * 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.celeborn.tests.flink
+
+import org.apache.flink.api.common.{ExecutionMode, RuntimeExecutionMode}
+import org.apache.flink.configuration.{Configuration, ExecutionOptions, RestOptions}
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.service.deploy.MiniClusterFeature
+import org.apache.celeborn.service.deploy.worker.Worker
+
+class SplitTest extends AnyFunSuite with Logging with MiniClusterFeature
+  with BeforeAndAfterAll {
+  var workers: collection.Set[Worker] = null
+  override def beforeAll(): Unit = {
+    logInfo("test initialized , setup celeborn mini cluster")
+    val masterConf = Map(
+      "celeborn.master.host" -> "localhost",
+      "celeborn.master.port" -> "9097")
+    val workerConf = Map(
+      "celeborn.master.endpoints" -> "localhost:9097",
+      CelebornConf.WORKER_FLUSHER_BUFFER_SIZE.key -> "10k")
+    workers = setUpMiniCluster(masterConf, workerConf)._2
+  }
+  override def afterAll(): Unit = {
+    logInfo("all test complete , stop celeborn mini cluster")
+    shutdownMiniCluster()
+  }
+
+  test("celeborn flink integration test - shuffle partition split test") {
+    val configuration = new Configuration
+    val parallelism = 8
+    configuration.setString(
+      "shuffle-service-factory.class",
+      "org.apache.celeborn.plugin.flink.RemoteShuffleServiceFactory")
+    configuration.setString(CelebornConf.MASTER_ENDPOINTS.key, "localhost:9097")
+    configuration.setString("execution.batch-shuffle-mode", "ALL_EXCHANGES_BLOCKING")
+    configuration.set(ExecutionOptions.RUNTIME_MODE, RuntimeExecutionMode.BATCH)
+    configuration.setString("taskmanager.memory.network.min", "1024m")
+    configuration.setString(RestOptions.BIND_PORT, "8081-8089")
+    configuration.setString(
+      "execution.batch.adaptive.auto-parallelism.min-parallelism",
+      "" + parallelism)
+    configuration.setString(
+      "execution.batch.adaptive.auto-parallelism.max-parallelism",
+      "" + parallelism)
+    configuration.setString("restart-strategy.type", "fixed-delay")
+    configuration.setString("restart-strategy.fixed-delay.attempts", "50")
+    configuration.setString("restart-strategy.fixed-delay.delay", "5s")
+    configuration.setString(CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD.key, "10k")
+    configuration.setString(CelebornConf.CLIENT_SHUFFLE_MAPPARTITION_SPLIT_ENABLED.key, "true")
+    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(configuration)
+    env.getConfig.setExecutionMode(ExecutionMode.BATCH)
+    env.getConfig.setParallelism(parallelism)
+    SplitHelper.runSplitRead(env)
+    env.execute("split test")
+  }
+
+}
--- a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
+++ b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
@ -21,8 +21,8 @@ import java.io.File

 import scala.collection.JavaConverters._

-import org.apache.flink.api.common.{ExecutionMode, InputDependencyConstraint, RuntimeExecutionMode}
-import org.apache.flink.configuration.{ConfigConstants, Configuration, ExecutionOptions, RestOptions}
+import org.apache.flink.api.common.{ExecutionMode, RuntimeExecutionMode}
+import org.apache.flink.configuration.{Configuration, ExecutionOptions, RestOptions}
 import org.apache.flink.runtime.jobgraph.JobType
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment
 import org.apache.flink.streaming.api.graph.GlobalStreamExchangeMode
@ -66,6 +66,9 @@ class WordCountTest extends AnyFunSuite with Logging with MiniClusterFeature
    configuration.setString(
      "execution.batch.adaptive.auto-parallelism.min-parallelism",
      "" + parallelism)
+    configuration.setString("restart-strategy.type", "fixed-delay")
+    configuration.setString("restart-strategy.fixed-delay.attempts", "50")
+    configuration.setString("restart-strategy.fixed-delay.delay", "5s")
    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(configuration)
    env.getConfig.setExecutionMode(ExecutionMode.BATCH)
    env.getConfig.setParallelism(parallelism)
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManager.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManager.java
@ -73,7 +73,7 @@ public class CreditStreamManager {
  }

  public long registerStream(
-      Consumer<Long> callback,
+      Consumer<Long> notifyStreamHandlerCallback,
      Channel channel,
      int initialCredit,
      int startSubIndex,
@ -117,7 +117,7 @@ public class CreditStreamManager {
    }
    mapDataPartition.tryRequestBufferOrRead();

-    callback.accept(streamId);
+    notifyStreamHandlerCallback.accept(streamId);
    addCredit(initialCredit, streamId);

    logger.debug("Register stream streamId: {}, fileInfo: {}", streamId, fileInfo);
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
@ -79,7 +79,6 @@ class MapDataPartition implements MemoryManager.ReadBufferTargetChangeListener {
    this.maxReadBuffers = maxReadBuffers;

    updateBuffersTarget((this.minReadBuffers + this.maxReadBuffers) / 2 + 1);
-
    logger.debug(
        "read map partition {} with {} {}",
        fileInfo.getFilePath(),
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartitionReader.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartitionReader.java
@ -37,6 +37,7 @@ import org.slf4j.LoggerFactory;
 import org.apache.celeborn.common.exception.FileCorruptedException;
 import org.apache.celeborn.common.meta.FileInfo;
 import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
 import org.apache.celeborn.common.network.protocol.ReadData;
 import org.apache.celeborn.common.network.protocol.TransportableError;
 import org.apache.celeborn.common.network.util.NettyUtils;
@ -436,7 +437,11 @@ public class MapDataPartitionReader implements Comparable<MapDataPartitionReader
    // we can safely release if reader reaches error or (read/send finished)
    synchronized (lock) {
      if (!isReleased) {
-        logger.debug("release reader for stream {}", this.streamId);
+        logger.debug("release reader for stream {}", streamId);
+        // old client can't support BufferStreamEnd, so for new client it tells client that this
+        // stream is finished.
+        if (fileInfo.isPartitionSplitEnabled() && !errorNotified)
+          associatedChannel.writeAndFlush(new BufferStreamEnd(streamId));
        if (!buffersToSend.isEmpty()) {
          numInUseBuffers.addAndGet(-1 * buffersToSend.size());
          buffersToSend.forEach(RecyclableBuffer::recycle);
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
@ -52,6 +52,7 @@ public final class MapPartitionFileWriter extends FileWriter {
  private long totalBytes;
  private long regionStartingOffset;
  private FileChannel indexChannel;
+  private volatile boolean isRegionFinished = true;

  public MapPartitionFileWriter(
      FileInfo fileInfo,
@ -120,8 +121,8 @@ public final class MapPartitionFileWriter extends FileWriter {
    long length = data.readableBytes();
    totalBytes += length;
    numSubpartitionBytes[partitionId] += length;
-
    super.write(data);
+    isRegionFinished = false;
  }

  @Override
@ -235,6 +236,7 @@ public final class MapPartitionFileWriter extends FileWriter {

    regionStartingOffset = totalBytes;
    Arrays.fill(numSubpartitionBytes, 0);
+    isRegionFinished = true;
  }

  private synchronized void destroyIndex() {
@ -301,4 +303,8 @@ public final class MapPartitionFileWriter extends FileWriter {

    return buffer;
  }
+
+  public boolean isRegionFinished() {
+    return isRegionFinished;
+  }
 }
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
@ -38,7 +38,7 @@ import org.apache.celeborn.common.protocol.message.ControlMessages._
 import org.apache.celeborn.common.protocol.message.StatusCode
 import org.apache.celeborn.common.rpc._
 import org.apache.celeborn.common.util.{JavaUtils, Utils}
-import org.apache.celeborn.service.deploy.worker.storage.StorageManager
+import org.apache.celeborn.service.deploy.worker.storage.{FileWriter, MapPartitionFileWriter, StorageManager}

 private[deploy] class Controller(
    override val rpcEnv: RpcEnv,
@ -90,7 +90,8 @@ private[deploy] class Controller(
          partitionType,
          rangeReadFilter,
          userIdentifier,
-          pushDataTimeout) =>
+          pushDataTimeout,
+          partitionSplitEnabled) =>
      val shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId)
      workerSource.sample(WorkerSource.RESERVE_SLOTS_TIME, shuffleKey) {
        logDebug(s"Received ReserveSlots request, $shuffleKey, " +
@ -107,7 +108,8 @@ private[deploy] class Controller(
          partitionType,
          rangeReadFilter,
          userIdentifier,
-          pushDataTimeout)
+          pushDataTimeout,
+          partitionSplitEnabled)
        logDebug(s"ReserveSlots for $shuffleKey finished.")
      }

@ -136,7 +138,8 @@ private[deploy] class Controller(
      partitionType: PartitionType,
      rangeReadFilter: Boolean,
      userIdentifier: UserIdentifier,
-      pushDataTimeout: Long): Unit = {
+      pushDataTimeout: Long,
+      partitionSplitEnabled: Boolean): Unit = {
    val shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId)
    if (shutdown.get()) {
      val msg = "Current worker is shutting down!"
@ -167,7 +170,8 @@ private[deploy] class Controller(
            splitMode,
            partitionType,
            rangeReadFilter,
-            userIdentifier)
+            userIdentifier,
+            partitionSplitEnabled)
          primaryLocs.add(new WorkingPartition(location, writer))
        } else {
          primaryLocs.add(location)
@ -206,7 +210,8 @@ private[deploy] class Controller(
            splitMode,
            partitionType,
            rangeReadFilter,
-            userIdentifier)
+            userIdentifier,
+            partitionSplitEnabled)
          replicaLocs.add(new WorkingPartition(location, writer))
        } else {
          replicaLocs.add(location)
@ -283,6 +288,7 @@ private[deploy] class Controller(
                }

                val fileWriter = location.asInstanceOf[WorkingPartition].getFileWriter
+                waitMapPartitionRegionFinished(fileWriter, conf.workerShuffleCommitTimeout)
                val bytes = fileWriter.close()
                if (bytes > 0L) {
                  if (fileWriter.getStorageInfo == null) {
@ -321,6 +327,23 @@ private[deploy] class Controller(
    future
  }

+  private def waitMapPartitionRegionFinished(fileWriter: FileWriter, waitTimeout: Long): Unit = {
+    if (fileWriter.isInstanceOf[MapPartitionFileWriter]) {
+      val delta = 100
+      var times = 0
+      while (delta * times < waitTimeout) {
+        if (fileWriter.asInstanceOf[MapPartitionFileWriter].isRegionFinished) {
+          logDebug(s"CommitFile succeed to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
+          return
+        }
+        Thread.sleep(delta)
+        times += 1
+      }
+      logWarning(
+        s"CommitFile faield to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
+    }
+  }
+
  private def handleCommitFiles(
      context: RpcCallContext,
      shuffleKey: String,
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@ -803,14 +803,6 @@ class PushDataHandler extends BaseMessageHandler with Logging {
        callback,
        wrappedCallback)) return

-    // During worker shutdown, worker will return HARD_SPLIT for all existed partition.
-    // This should before return exception to make current push request revive and retry.
-    if (shutdown.get()) {
-      logInfo(s"Push data return HARD_SPLIT for shuffle $shuffleKey since worker shutdown.")
-      callback.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
-      return
-    }
-
    val fileWriter =
      getFileWriterAndCheck(pushData.`type`(), location, isPrimary, callback) match {
        case (true, _) => return
@ -860,7 +852,7 @@ class PushDataHandler extends BaseMessageHandler with Logging {
    val msg = Message.decode(rpcRequest.body().nioByteBuffer())
    val requestId = rpcRequest.requestId
    val (mode, shuffleKey, partitionUniqueId, checkSplit) = msg match {
-      case p: PushDataHandShake => (p.mode, p.shuffleKey, p.partitionUniqueId, false)
+      case p: PushDataHandShake => (p.mode, p.shuffleKey, p.partitionUniqueId, true)
      case rs: RegionStart => (rs.mode, rs.shuffleKey, rs.partitionUniqueId, true)
      case rf: RegionFinish => (rf.mode, rf.shuffleKey, rf.partitionUniqueId, false)
    }
@ -869,7 +861,7 @@ class PushDataHandler extends BaseMessageHandler with Logging {
      rpcRequest,
      requestId,
      () =>
-        handleRpcRequestCore(
+        handleMapPartitionRpcRequestCore(
          mode,
          msg,
          shuffleKey,
@ -883,7 +875,7 @@ class PushDataHandler extends BaseMessageHandler with Logging {

  }

-  private def handleRpcRequestCore(
+  private def handleMapPartitionRpcRequestCore(
      mode: Byte,
      message: Message,
      shuffleKey: String,
@ -943,7 +935,22 @@ class PushDataHandler extends BaseMessageHandler with Logging {
        case (false, f: FileWriter) => f
      }

-    if (checkSplit && checkDiskFullAndSplit(fileWriter, isPrimary, null, callback)) return
+    // During worker shutdown, worker will return HARD_SPLIT for all existed partition.
+    // This should before return exception to make current push request revive and retry.
+    val isPartitionSplitEnabled = fileWriter.asInstanceOf[
+      MapPartitionFileWriter].getFileInfo.isPartitionSplitEnabled
+
+    if (shutdown.get() && (messageType == Type.REGION_START || messageType == Type.PUSH_DATA_HAND_SHAKE) && isPartitionSplitEnabled) {
+      logInfo(s"$messageType return HARD_SPLIT for shuffle $shuffleKey since worker shutdown.")
+      callback.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
+      return
+    }
+
+    if (checkSplit && (messageType == Type.REGION_START || messageType == Type.PUSH_DATA_HAND_SHAKE) && isPartitionSplitEnabled && checkDiskFullAndSplit(
+        fileWriter,
+        isPrimary,
+        null,
+        callback)) return

    try {
      messageType match {
@ -1108,6 +1115,8 @@ class PushDataHandler extends BaseMessageHandler with Logging {
      softSplit: AtomicBoolean,
      callback: RpcResponseCallback): Boolean = {
    val diskFull = checkDiskFull(fileWriter)
+    logDebug(
+      s"CheckDiskFullAndSplit in diskfull: $diskFull, partitionSplitMinimumSize: $partitionSplitMinimumSize, splitThreshold: ${fileWriter.getSplitThreshold()}, filelength: ${fileWriter.getFileInfo.getFileLength}, filename:${fileWriter.getFileInfo.getFilePath}")
    if (workerPartitionSplitEnabled && ((diskFull && fileWriter.getFileInfo.getFileLength > partitionSplitMinimumSize) ||
        (isPrimary && fileWriter.getFileInfo.getFileLength > fileWriter.getSplitThreshold()))) {
      if (softSplit != null && fileWriter.getSplitMode == PartitionSplitMode.SOFT &&
@ -1115,6 +1124,8 @@ class PushDataHandler extends BaseMessageHandler with Logging {
        softSplit.set(true)
      } else {
        callback.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
+        logDebug(
+          s"CheckDiskFullAndSplit hardsplit diskfull: $diskFull, partitionSplitMinimumSize: $partitionSplitMinimumSize, splitThreshold: ${fileWriter.getSplitThreshold()}, filelength: ${fileWriter.getFileInfo.getFileLength}, filename:${fileWriter.getFileInfo.getFilePath}")
        return true
      }
    }
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@ -301,6 +301,29 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
      partitionType: PartitionType,
      rangeReadFilter: Boolean,
      userIdentifier: UserIdentifier): FileWriter = {
+    createWriter(
+      appId,
+      shuffleId,
+      location,
+      splitThreshold,
+      splitMode,
+      partitionType,
+      rangeReadFilter,
+      userIdentifier,
+      true)
+  }
+
+  @throws[IOException]
+  def createWriter(
+      appId: String,
+      shuffleId: Int,
+      location: PartitionLocation,
+      splitThreshold: Long,
+      splitMode: PartitionSplitMode,
+      partitionType: PartitionType,
+      rangeReadFilter: Boolean,
+      userIdentifier: UserIdentifier,
+      partitionSplitEnabled: Boolean): FileWriter = {
    if (healthyWorkingDirs().size <= 0 && !hasHDFSStorage) {
      throw new IOException("No available working dirs!")
    }
@ -328,7 +351,11 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
          new Path(new Path(hdfsDir, conf.workerWorkingDir), s"$appId/$shuffleId")
        FileSystem.mkdirs(StorageManager.hadoopFs, shuffleDir, hdfsPermission)
        val fileInfo =
-          new FileInfo(new Path(shuffleDir, fileName).toString, userIdentifier, partitionType)
+          new FileInfo(
+            new Path(shuffleDir, fileName).toString,
+            userIdentifier,
+            partitionType,
+            partitionSplitEnabled)
        val hdfsWriter = partitionType match {
          case PartitionType.MAP => new MapPartitionFileWriter(
              fileInfo,
@ -374,7 +401,12 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
                s"Create shuffle data file ${file.getAbsolutePath} failed!")
            }
          }
-          val fileInfo = new FileInfo(file.getAbsolutePath, userIdentifier, partitionType)
+          val fileInfo =
+            new FileInfo(
+              file.getAbsolutePath,
+              userIdentifier,
+              partitionType,
+              partitionSplitEnabled)
          fileInfo.setMountPoint(mountPoint)
          val fileWriter = partitionType match {
            case PartitionType.MAP => new MapPartitionFileWriter(
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java
@ -101,9 +101,8 @@ public class CreditStreamManagerSuiteJ {

    mapDataPartition1.getStreamReader(registerStream1).recycle();

-    timeOutOrMeetCondition(() -> creditStreamManager.numRecycleStreams() == 0);
+    timeOutOrMeetCondition(() -> creditStreamManager.numStreamStates() == 3);
    Assert.assertEquals(creditStreamManager.numRecycleStreams(), 0);
-    Assert.assertEquals(3, creditStreamManager.numStreamStates());

    // registerStream2 can't be cleaned as registerStream2 is not finished
    AtomicInteger numInFlightRequests =
@ -117,8 +116,10 @@ public class CreditStreamManagerSuiteJ {
    // recycle all channel
    numInFlightRequests.decrementAndGet();
    creditStreamManager.connectionTerminated(channel);
-    timeOutOrMeetCondition(() -> creditStreamManager.numRecycleStreams() == 0);
-    Assert.assertEquals(creditStreamManager.numStreamStates(), 0);
+    timeOutOrMeetCondition(() -> creditStreamManager.numStreamStates() == 0);
+    // when cpu is busy, even through that timeOutOrMeetCondition is true,
+    // creditStreamManager.numStreamStates are still not be removed
+    Assert.assertTrue(creditStreamManager.numRecycleStreams() >= 0);
  }

  @AfterClass
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
@ -20,6 +20,9 @@ package org.apache.celeborn.service.deploy.cluster
 import java.io.ByteArrayOutputStream
 import java.nio.charset.StandardCharsets

+import scala.collection.mutable
+import scala.util.control.Breaks
+
 import org.apache.commons.lang3.RandomStringUtils
 import org.junit.Assert
 import org.scalatest.BeforeAndAfterAll
@ -64,8 +67,10 @@ trait ReadWriteTestBase extends AnyFunSuite
    val lifecycleManager = new LifecycleManager(APP, clientConf)
    val shuffleClient = new ShuffleClientImpl(APP, clientConf, UserIdentifier("mock", "mock"))
    shuffleClient.setupLifecycleManagerRef(lifecycleManager.self)
-
-    val STR1 = RandomStringUtils.random(1024)
+    val dataPrefix = Array("000000", "111111", "222222", "333333")
+    val dataPrefixMap = new mutable.HashMap[String, String]
+    val STR1 = dataPrefix(0) + RandomStringUtils.random(1024)
+    dataPrefixMap.put(dataPrefix(0), STR1)
    val DATA1 = STR1.getBytes(StandardCharsets.UTF_8)
    val OFFSET1 = 0
    val LENGTH1 = DATA1.length
@ -73,19 +78,22 @@ trait ReadWriteTestBase extends AnyFunSuite
    val dataSize1 = shuffleClient.pushData(1, 0, 0, 0, DATA1, OFFSET1, LENGTH1, 1, 1)
    logInfo(s"push data data size $dataSize1")

-    val STR2 = RandomStringUtils.random(32 * 1024)
+    val STR2 = dataPrefix(1) + RandomStringUtils.random(32 * 1024)
+    dataPrefixMap.put(dataPrefix(1), STR2)
    val DATA2 = STR2.getBytes(StandardCharsets.UTF_8)
    val OFFSET2 = 0
    val LENGTH2 = DATA2.length
    val dataSize2 = shuffleClient.pushData(1, 0, 0, 0, DATA2, OFFSET2, LENGTH2, 1, 1)
    logInfo("push data data size " + dataSize2)

-    val STR3 = RandomStringUtils.random(32 * 1024)
+    val STR3 = dataPrefix(2) + RandomStringUtils.random(32 * 1024)
+    dataPrefixMap.put(dataPrefix(2), STR3)
    val DATA3 = STR3.getBytes(StandardCharsets.UTF_8)
    val LENGTH3 = DATA3.length
    shuffleClient.mergeData(1, 0, 0, 0, DATA3, 0, LENGTH3, 1, 1)

-    val STR4 = RandomStringUtils.random(16 * 1024)
+    val STR4 = dataPrefix(3) + RandomStringUtils.random(16 * 1024)
+    dataPrefixMap.put(dataPrefix(3), STR4)
    val DATA4 = STR4.getBytes(StandardCharsets.UTF_8)
    val LENGTH4 = DATA4.length
    shuffleClient.mergeData(1, 0, 0, 0, DATA4, 0, LENGTH4, 1, 1)
@ -104,9 +112,12 @@ trait ReadWriteTestBase extends AnyFunSuite
    }

    val readBytes = outputStream.toByteArray
+    val readStringMap = getReadStringMap(readBytes, dataPrefix, dataPrefixMap)
+
    Assert.assertEquals(LENGTH1 + LENGTH2 + LENGTH3 + LENGTH4, readBytes.length)
-    val targetArr = Array.concat(DATA1, DATA2, DATA3, DATA4)
-    Assert.assertArrayEquals(targetArr, readBytes)
+    for ((prefix, data) <- readStringMap) {
+      Assert.assertEquals(dataPrefixMap.get(prefix).get, data)
+    }

    Thread.sleep(5000L)
    shuffleClient.shutdown()
@ -114,4 +125,28 @@ trait ReadWriteTestBase extends AnyFunSuite

  }

+  def getReadStringMap(
+      readBytes: Array[Byte],
+      dataPrefix: Array[String],
+      dataPrefixMap: mutable.HashMap[String, String]): mutable.HashMap[String, String] = {
+    var readString = new String(readBytes, StandardCharsets.UTF_8)
+    val prefixStringMap = new mutable.HashMap[String, String]
+    val loop = new Breaks;
+    for (i <- 0 to 4) {
+      loop.breakable {
+        for (prefix <- dataPrefix) {
+          if (readString.startsWith(prefix)) {
+            val subString = readString.substring(0, dataPrefixMap.get(prefix).get.length)
+            prefixStringMap.put(prefix, subString)
+            println(
+              s"readString before: ${readString.length}, ${dataPrefixMap.get(prefix).get.length}")
+            readString = readString.substring(dataPrefixMap.get(prefix).get.length)
+            println(s"readString after: ${readString.length}")
+            loop.break()
+          }
+        }
+      }
+    }
+    prefixStringMap
+  }
 }